ob-metaflow 2.12.10.1__py2.py3-none-any.whl → 2.12.10.1rc2__py2.py3-none-any.whl

metaflow/plugins/datastores/gs_storage.py

@@ -8,6 +8,7 @@ from tempfile import mkdtemp
 
 from metaflow.datastore.datastore_storage import DataStoreStorage, CloseAfterUse
 from metaflow.exception import MetaflowInternalError
+from metaflow.tracing import traced, tracing
 from metaflow.metaflow_config import (
     DATASTORE_SYSROOT_GS,
     ARTIFACT_LOCALROOT,
@@ -70,12 +71,14 @@ class _GSRootClient(object):
         """Drives GSStorage.is_file()"""
         try:
             blob = self.get_blob_client(path)
-            result = blob.exists()
+            with traced("exists", dict(path=path)):
+                result = blob.exists()
 
             return result
         except Exception as e:
             process_gs_exception(e)
 
+    @tracing
     def list_content_single(self, path):
         """Drives GSStorage.list_content()"""
 
@@ -104,6 +107,7 @@ class _GSRootClient(object):
         except Exception as e:
             process_gs_exception(e)
 
+    @tracing
     def save_bytes_single(
         self,
         path_tmpfile_metadata_triple,
@@ -119,10 +123,12 @@ class _GSRootClient(object):
                 blob.metadata = {"metaflow-user-attributes": json.dumps(metadata)}
             from google.cloud.storage.retry import DEFAULT_RETRY
 
-            blob.upload_from_filename(tmpfile, retry=DEFAULT_RETRY)
+            with traced("upload_blob", dict(path=path)):
+                blob.upload_from_filename(tmpfile, retry=DEFAULT_RETRY)
         except Exception as e:
             process_gs_exception(e)
 
+    @tracing
     def load_bytes_single(self, tmpdir, key):
         """Drives GSStorage.load_bytes()"""
         tmp_filename = os.path.join(tmpdir, str(uuid.uuid4()))
@@ -136,7 +142,8 @@ class _GSRootClient(object):
                 metaflow_user_attributes = json.loads(
                     blob.metadata["metaflow-user-attributes"]
                 )
-            blob.download_to_filename(tmp_filename)
+            with traced("download_blob", dict(path=key)):
+                blob.download_to_filename(tmp_filename)
         except google.api_core.exceptions.NotFound:
             tmp_filename = None
         return key, tmp_filename, metaflow_user_attributes

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -5,6 +5,7 @@ import random
 import sys
 import time
 
+from metaflow.tracing import inject_tracing_vars, tracing, traced
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import KUBERNETES_SECRETS
 from metaflow.tracing import inject_tracing_vars
@@ -299,13 +300,19 @@ class KubernetesJob(object):
             #       achieve the guarantees that we are seeking.
             #       https://github.com/kubernetes/enhancements/issues/1040
             #       Hopefully, we will be able to get creative with kube-batch
-            response = (
-                client.BatchV1Api()
-                .create_namespaced_job(
-                    body=self._job, namespace=self._kwargs["namespace"]
+
+            with traced(
+                "submit_kubernetes_job",
+                {"job": "%s/%s" % (self._job.metadata.name, self._kwargs["namespace"])},
+            ):
+
+                response = (
+                    client.BatchV1Api()
+                    .create_namespaced_job(
+                        body=self._job, namespace=self._kwargs["namespace"]
+                    )
+                    .to_dict()
                 )
-                .to_dict()
-            )
             return RunningJob(
                 client=self._client,
                 name=response["metadata"]["name"],
@@ -434,11 +441,14 @@ class RunningJob(object):
     def _fetch_job(self):
         client = self._client.get()
         try:
-            return (
-                client.BatchV1Api()
-                .read_namespaced_job(name=self._name, namespace=self._namespace)
-                .to_dict()
-            )
+            with traced(
+                "fetch_kubernetes_job", {"job": "%s/%s" % (self._name, self._namespace)}
+            ):
+                return (
+                    client.BatchV1Api()
+                    .read_namespaced_job(name=self._name, namespace=self._namespace)
+                    .to_dict()
+                )
         except client.rest.ApiException as e:
             if e.status == 404:
                 raise KubernetesJobException(
@@ -450,17 +460,20 @@ class RunningJob(object):
     def _fetch_pod(self):
         # Fetch pod metadata.
         client = self._client.get()
-        pods = (
-            client.CoreV1Api()
-            .list_namespaced_pod(
-                namespace=self._namespace,
-                label_selector="job-name={}".format(self._name),
+        with traced(
+            "fetch_kubernetes_pod", {"job": "%s/%s" % (self._name, self._namespace)}
+        ):
+            pods = (
+                client.CoreV1Api()
+                .list_namespaced_pod(
+                    namespace=self._namespace,
+                    label_selector="job-name={}".format(self._name),
+                )
+                .to_dict()["items"]
             )
-            .to_dict()["items"]
-        )
-        if pods:
-            return pods[0]
-        return {}
+            if pods:
+                return pods[0]
+            return {}
 
     def kill(self):
         # Terminating a Kubernetes job is a bit tricky. Issuing a

metaflow/plugins/storage_executor.py

@@ -6,6 +6,7 @@ import sys
 from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
 
 from metaflow.exception import MetaflowException
+from metaflow.tracing import TracedThreadPoolExecutor
 
 if sys.version_info[:2] < (3, 7):
     # in 3.6, Only BrokenProcessPool exists (there is no BrokenThreadPool)
@@ -132,7 +133,10 @@ class StorageExecutor(object):
                     msg="Cannot use ProcessPoolExecutor because Python version is older than 3.7 and multiprocess start method has been set to something other than 'spawn'"
                 )
         else:
-            self._executor = ThreadPoolExecutor(max_workers=threadpool_max_workers)
+            self._executor = TracedThreadPoolExecutor(
+                "StorageExecutor", max_workers=threadpool_max_workers
+            )
+            # self._executor = ThreadPoolExecutor(max_workers=threadpool_max_workers)
 
     def warm_up(self):
         # warm up at least one process or thread in the pool.

metaflow/tracing/__init__.py

@@ -52,6 +52,10 @@ def tracing(func):
     return wrapper_func
 
 
+from concurrent.futures import ThreadPoolExecutor
+TracedThreadPoolExecutor = ThreadPoolExecutor
+
+
 if not DISABLE_TRACING and (CONSOLE_TRACE_ENABLED or OTEL_ENDPOINT or ZIPKIN_ENDPOINT):
     try:
         # Overrides No-Op implementations if a specific provider is configured.
@@ -63,6 +67,7 @@ if not DISABLE_TRACING and (CONSOLE_TRACE_ENABLED or OTEL_ENDPOINT or ZIPKIN_END
             get_trace_id,
             traced,
             tracing,
+            TracedThreadPoolExecutor,
         )
 
     except ImportError as e:

metaflow/tracing/threadpool.py

@@ -0,0 +1,30 @@
+from concurrent.futures import ThreadPoolExecutor
+from opentelemetry import context as otel_context
+from opentelemetry.sdk.trace import Tracer
+from opentelemetry import trace as trace_api
+from typing import Callable
+
+
+class TracedThreadPoolExecutor(ThreadPoolExecutor):
+    """Implementation of :class:`ThreadPoolExecutor` that will pass context into sub tasks."""
+
+    def __init__(self, name: str, *args, **kwargs):
+        tracer = trace_api.get_tracer(name)
+        self.tracer = tracer
+        super().__init__(*args, **kwargs)
+
+    def with_otel_context(self, context: otel_context.Context, fn: Callable):
+        otel_context.attach(context)
+        return fn()
+
+    def submit(self, fn, *args, **kwargs):
+        """Submit a new task to the thread pool."""
+
+        # get the current otel context
+        context = otel_context.get_current()
+        if context:
+            return super().submit(
+                lambda: self.with_otel_context(context, lambda: fn(*args, **kwargs)),
+            )
+        else:
+            return super().submit(lambda: fn(*args, **kwargs))

metaflow/tracing/tracing_modules.py

@@ -13,6 +13,8 @@ from typing import Dict, List, Optional
 from opentelemetry import trace as trace_api, context
 from .span_exporter import get_span_exporter
 
+from .threadpool import TracedThreadPoolExecutor
+
 tracer_provider = None
 
 
@@ -21,6 +23,7 @@ def init_tracing():
     if tracer_provider is not None:
         print("Tracing already initialized", file=sys.stderr)
         return
+    import os
 
     from .propagator import EnvPropagator
 
@@ -142,7 +145,7 @@ def traced(name, attrs={}):
 def tracing(func):
     @wraps(func)
     def wrapper_func(*args, **kwargs):
-        tracer = trace_api.get_tracer_provider().get_tracer(func.__module__.__name__)
+        tracer = trace_api.get_tracer_provider().get_tracer(func.__module__)
 
         with tracer.start_as_current_span(func.__name__):
             return func(*args, **kwargs)

metaflow/version.py

@@ -1 +1 @@
-metaflow_version = "2.12.10.1"
+metaflow_version = "2.12.10.1rc2"

{ob_metaflow-2.12.10.1.dist-info → ob_metaflow-2.12.10.1rc2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow
-Version: 2.12.10.1
+Version: 2.12.10.1rc2
 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.12.10.1 ; extra == 'stubs'
+Requires-Dist: metaflow-stubs ==2.12.10.1rc2 ; extra == 'stubs'
 
 ![Metaflow_Logo_Horizontal_FullColor_Ribbon_Dark_RGB](https://user-images.githubusercontent.com/763451/89453116-96a57e00-d713-11ea-9fa6-82b29d4d6eff.png)
 

{ob_metaflow-2.12.10.1.dist-info → ob_metaflow-2.12.10.1rc2.dist-info}/RECORD

@@ -35,7 +35,7 @@ metaflow/tuple_util.py,sha256=_G5YIEhuugwJ_f6rrZoelMFak3DqAR2tt_5CapS1XTY,830
 metaflow/unbounded_foreach.py,sha256=p184WMbrMJ3xKYHwewj27ZhRUsSj_kw1jlye5gA9xJk,387
 metaflow/util.py,sha256=m5womQ7y-jXehuMyHPfByDbZ4HwTJxzs869cPOlMR8s,13057
 metaflow/vendor.py,sha256=FchtA9tH22JM-eEtJ2c9FpUdMn8sSb1VHuQS56EcdZk,5139
-metaflow/version.py,sha256=3IvPc0z5CHdysfIGeD_OJlA1QS8xf5sReJcs7eZ-NRU,31
+metaflow/version.py,sha256=kGj-PcLLHyv0B6fbtr3neBGsNkLetAXjMRmb2CychSs,34
 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
@@ -154,7 +154,7 @@ metaflow/plugins/parallel_decorator.py,sha256=wXBmEnnOmWI46-Lcry_CArkujj6YOrfzAw
 metaflow/plugins/project_decorator.py,sha256=eJOe0Ea7CbUCReEhR_XQvRkhV6jyRqDxM72oZI7EMCk,5336
 metaflow/plugins/resources_decorator.py,sha256=3MMZ7uptDf99795_RcSOq4h0N3OFlKpd3ahIEsozBBs,1333
 metaflow/plugins/retry_decorator.py,sha256=tz_2Tq6GLg3vjDBZp0KKVTk3ADlCvqaWTSf7blmFdUw,1548
-metaflow/plugins/storage_executor.py,sha256=FqAgR0-L9MuqN8fRtTe4jjUfJL9lqt6fQkYaglAjRbk,6137
+metaflow/plugins/storage_executor.py,sha256=BUzSyv0cU5I_FN8RfXpNy6mlTw_8UKa-xm_5AUJ8aGQ,6332
 metaflow/plugins/tag_cli.py,sha256=O_ZI4ILwGX3xKrLewUUF-zdJjCDi3JmsTb4ow87_RuY,17610
 metaflow/plugins/test_unbounded_foreach_decorator.py,sha256=-3o_VhmZW7R9i-0RgqEradmPqx9rS6jJxcIKV6-WQMg,5948
 metaflow/plugins/timeout_decorator.py,sha256=R-X8rKeMqd-xhfJFqskWb6ZpmZt2JB14U1BZJSRriwM,3648
@@ -237,7 +237,7 @@ metaflow/plugins/cards/card_modules/chevron/tokenizer.py,sha256=lQU9OELUE9a5Xu4s
 metaflow/plugins/cards/card_viewer/viewer.html,sha256=qZJGzhZhQ1gugsknRP7zkAPPfUAtvemK1UKqXoGff5M,11593
 metaflow/plugins/datastores/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow/plugins/datastores/azure_storage.py,sha256=QvIUQGOZF1oKeRJXbl3RsV1MnO3OMfLzheoo0UNWn7E,16670
-metaflow/plugins/datastores/gs_storage.py,sha256=ED1xfXNDiPkqma5PEui8HCv0b4ImHK_rlUMQJD9UNes,9622
+metaflow/plugins/datastores/gs_storage.py,sha256=qKdYuMLw2_PQRRWWEWu57MymRAsgeTe6dY5Lmprh_lw,9885
 metaflow/plugins/datastores/local_storage.py,sha256=igrBDphhyu7EFIUj3BWcO7beiZbNnJLq--lF45UYSyI,4750
 metaflow/plugins/datastores/s3_storage.py,sha256=CZdNqaKtxDXQbEg2YHyphph3hWcLIE50puenm0WGVpk,5473
 metaflow/plugins/datatools/__init__.py,sha256=ge4L16OBQLy2J_MMvoHg3lMfdm-MluQgRWoyZ5GCRnk,1267
@@ -281,7 +281,7 @@ metaflow/plugins/kubernetes/kubernetes.py,sha256=0UjKZy5_5KOGYC-k0DZqpv-ca_tZBkw
 metaflow/plugins/kubernetes/kubernetes_cli.py,sha256=qBDdr1Lvtt-RO9pB-9_HTOPdzAmDvvJ0aiQ1OoCcrMU,10892
 metaflow/plugins/kubernetes/kubernetes_client.py,sha256=GKg-gT3qhXMRQV-sG1YyoOf3Z32NXr_wwEN2ytMVSEg,2471
 metaflow/plugins/kubernetes/kubernetes_decorator.py,sha256=Xu1PB9hrodnDcPIVHMC9HrTvOQzsEr4nR80zoIlXTME,25463
-metaflow/plugins/kubernetes/kubernetes_job.py,sha256=1uQmEHsLyJ5jewmIA5smrFtSO8MymOqPDb-YIS65u50,31614
+metaflow/plugins/kubernetes/kubernetes_job.py,sha256=VcRhUrAlPQMwJ02xGHUh-E4dh1EWP-6RjwjgIiZHDf8,32190
 metaflow/plugins/kubernetes/kubernetes_jobsets.py,sha256=Cr5M7m0lKShv6o8xjRJap8hSRgf51fz4RR0THlCSM08,40780
 metaflow/plugins/metadata/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 metaflow/plugins/metadata/local.py,sha256=YhLJC5zjVJrvQFIyQ92ZBByiUmhCC762RUX7ITX12O8,22428
@@ -315,10 +315,11 @@ metaflow/system/__init__.py,sha256=SB9Py7Acecqi76MY9MonSHXFuDD1yIJEGJtEQH8cNq4,1
 metaflow/system/system_logger.py,sha256=31noRo2qFdFyVtlyvIV2coUw4x6YdJGCdhctNBbBUso,3299
 metaflow/system/system_monitor.py,sha256=NracI4ITbUZWziGsAvE1910ODZhXGVYro7jD-rLWe48,3578
 metaflow/system/system_utils.py,sha256=E5C66_oeAxEE4-2okKlA-X6rgeYndPFP7KmWKHvvFZ8,657
-metaflow/tracing/__init__.py,sha256=xYTOT5BS5jbwhjk6hskxqNSU9st2LYtfeLN2Hknm3EI,1551
+metaflow/tracing/__init__.py,sha256=aOoU4o0fIHARKY25QUeqrRpJMtXS9r2NcSvakPu047k,1687
 metaflow/tracing/propagator.py,sha256=AdPeAqoeRauH82pTw01hLFNPRAzm29nlwM7C2iqKFFk,2502
 metaflow/tracing/span_exporter.py,sha256=eFVW40Fakct26KWycA-HEyl9yDqZRSJoLr3BOx8qopM,1697
-metaflow/tracing/tracing_modules.py,sha256=_-3xwEq2-qRY2o2QMtZI6Xl8G5OHv_SAJUNaan0EYt8,4877
+metaflow/tracing/threadpool.py,sha256=ldoSO9baoNPJpfr9msclNlQeeRfJIysGQ48YqLniM1k,1082
+metaflow/tracing/tracing_modules.py,sha256=fhuflMFkMLCZnGw3uccn22DMp87EE5oblogjQnbeZyg,4932
 metaflow/tutorials/00-helloworld/README.md,sha256=ML9k-2mRZOAFV3RDzwAWSIjnWzG03-3n9ov59UAvN5k,386
 metaflow/tutorials/00-helloworld/helloworld.py,sha256=ji0v7TFjVVR0VTs-shQOswh1Zd0YP5rZxud_LEmsjHU,857
 metaflow/tutorials/01-playlist/README.md,sha256=b9jgxen-vaVNIyH2UZPys4vvbxBylSoqDzCQ-aL3ves,838
@@ -342,9 +343,9 @@ metaflow/tutorials/07-worldview/README.md,sha256=5vQTrFqulJ7rWN6r20dhot9lI2sVj9W
 metaflow/tutorials/07-worldview/worldview.ipynb,sha256=ztPZPI9BXxvW1QdS2Tfe7LBuVzvFvv0AToDnsDJhLdE,2237
 metaflow/tutorials/08-autopilot/README.md,sha256=GnePFp_q76jPs991lMUqfIIh5zSorIeWznyiUxzeUVE,1039
 metaflow/tutorials/08-autopilot/autopilot.ipynb,sha256=DQoJlILV7Mq9vfPBGW-QV_kNhWPjS5n6SJLqePjFYLY,3191
-ob_metaflow-2.12.10.1.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
-ob_metaflow-2.12.10.1.dist-info/METADATA,sha256=RMXh3hmIQPNiE9VPNB4fcJfshMltyyQj8ECecFjkd9A,5145
-ob_metaflow-2.12.10.1.dist-info/WHEEL,sha256=XRxW4r1PNiVhMpP4bT9oWtu3HyndxpJ84SkubFgzp_Y,109
-ob_metaflow-2.12.10.1.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
-ob_metaflow-2.12.10.1.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
-ob_metaflow-2.12.10.1.dist-info/RECORD,,
+ob_metaflow-2.12.10.1rc2.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
+ob_metaflow-2.12.10.1rc2.dist-info/METADATA,sha256=ZDGv_s5Zo4rCdA5Ds0EJjYYxOVYoCBhwEnPwh-YApJ4,5151
+ob_metaflow-2.12.10.1rc2.dist-info/WHEEL,sha256=XRxW4r1PNiVhMpP4bT9oWtu3HyndxpJ84SkubFgzp_Y,109
+ob_metaflow-2.12.10.1rc2.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
+ob_metaflow-2.12.10.1rc2.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
+ob_metaflow-2.12.10.1rc2.dist-info/RECORD,,