ob-metaflow 2.12.9.1__tar.gz → 2.12.10.1rc1__tar.gz

{ob_metaflow-2.12.9.1/ob_metaflow.egg-info → ob_metaflow-2.12.10.1rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow
-Version: 2.12.9.1
+Version: 2.12.10.1rc1
 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.9.1; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.12.10.1rc1; 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.9.1 → ob_metaflow-2.12.10.1rc1}/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

{ob_metaflow-2.12.9.1 → ob_metaflow-2.12.10.1rc1}/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

{ob_metaflow-2.12.9.1 → ob_metaflow-2.12.10.1rc1}/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.threadpool 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.

ob_metaflow-2.12.10.1rc1/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))

{ob_metaflow-2.12.9.1 → ob_metaflow-2.12.10.1rc1}/metaflow/tracing/tracing_modules.py

@@ -21,6 +21,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 +143,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)

ob_metaflow-2.12.10.1rc1/metaflow/version.py

@@ -0,0 +1 @@
+metaflow_version = "2.12.10.1rc1"

{ob_metaflow-2.12.9.1 → ob_metaflow-2.12.10.1rc1/ob_metaflow.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow
-Version: 2.12.9.1
+Version: 2.12.10.1rc1
 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.9.1; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.12.10.1rc1; 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.9.1 → ob_metaflow-2.12.10.1rc1}/ob_metaflow.egg-info/SOURCES.txt

@@ -323,6 +323,7 @@ metaflow/system/system_utils.py
 metaflow/tracing/__init__.py
 metaflow/tracing/propagator.py
 metaflow/tracing/span_exporter.py
+metaflow/tracing/threadpool.py
 metaflow/tracing/tracing_modules.py
 metaflow/tutorials/00-helloworld/README.md
 metaflow/tutorials/00-helloworld/helloworld.py

{ob_metaflow-2.12.9.1 → ob_metaflow-2.12.10.1rc1}/ob_metaflow.egg-info/requires.txt

@@ -4,4 +4,4 @@ pylint
 kubernetes
 
 [stubs]
-metaflow-stubs==2.12.9.1
+metaflow-stubs==2.12.10.1rc1

ob_metaflow-2.12.9.1/metaflow/version.py

@@ -1 +0,0 @@
-metaflow_version = "2.12.9.1"