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

metaflow/tracing/__init__.py

@@ -20,15 +20,15 @@ def post_fork():
     yield
 
 
-def cli_entrypoint(name: str):
-    def cli_entrypoint_wrap(func):
+def cli(name: str):
+    def cli_wrap(func):
         @wraps(func)
         def wrapper_func(*args, **kwargs):
             return func(*args, **kwargs)
 
         return wrapper_func
 
-    return cli_entrypoint_wrap
+    return cli_wrap
 
 
 def inject_tracing_vars(env_dict: Dict[str, str]) -> Dict[str, str]:
@@ -40,7 +40,9 @@ def get_trace_id() -> str:
 
 
 @contextlib.contextmanager
-def traced(name, attrs={}):
+def traced(name, attrs=None):
+    if attrs is None:
+        attrs = {}
     yield
 
 
@@ -54,17 +56,15 @@ def tracing(func):
 
 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.
         from .tracing_modules import (
             init_tracing,
             post_fork,
-            cli_entrypoint,
+            cli,
             inject_tracing_vars,
             get_trace_id,
             traced,
             tracing,
         )
-
     except ImportError as e:
         # We keep the errors silent by default so that having tracing environment variables present
         # does not affect users with no need for tracing.

metaflow/tracing/span_exporter.py

@@ -3,6 +3,8 @@ from metaflow.metaflow_config import (
     OTEL_ENDPOINT,
     ZIPKIN_ENDPOINT,
     CONSOLE_TRACE_ENABLED,
+    SERVICE_AUTH_KEY,
+    SERVICE_HEADERS,
 )
 
 if OTEL_ENDPOINT:
@@ -11,56 +13,47 @@ if OTEL_ENDPOINT:
 
 
 def get_span_exporter():
-    if OTEL_ENDPOINT:
-        return set_otel_exporter()
+    exporter_map = {
+        OTEL_ENDPOINT: _create_otel_exporter,
+        ZIPKIN_ENDPOINT: _create_zipkin_exporter,
+        CONSOLE_TRACE_ENABLED: _create_console_exporter,
+    }
 
-    elif ZIPKIN_ENDPOINT:
-        return set_zipkin_exporter()
+    for config, create_exporter in exporter_map.items():
+        if config:
+            return create_exporter()
 
-    elif CONSOLE_TRACE_ENABLED:
-        return set_console_exporter()
-    else:
-        print("WARNING: endpoints not set up for Opentelemetry", file=sys.stderr)
-        return
+    print("WARNING: endpoints not set up for OpenTelemetry", file=sys.stderr)
+    return None
 
 
-def set_otel_exporter():
+def _create_otel_exporter():
     from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 
-    from metaflow.metaflow_config import (
-        SERVICE_AUTH_KEY,
-        SERVICE_HEADERS,
-    )
+    if not any([SERVICE_AUTH_KEY, SERVICE_HEADERS]):
+        print("WARNING: no auth settings for OpenTelemetry", file=sys.stderr)
+        return None
+
+    config = {
+        "endpoint": OTEL_ENDPOINT,
+        "timeout": 15,
+    }
 
     if SERVICE_AUTH_KEY:
-        span_exporter = OTLPSpanExporter(
-            endpoint=OTEL_ENDPOINT,
-            headers={"x-api-key": SERVICE_AUTH_KEY},
-            timeout=15,
-        )
+        config["headers"] = {"x-api-key": SERVICE_AUTH_KEY}
     elif SERVICE_HEADERS:
-        span_exporter = OTLPSpanExporter(
-            endpoint=OTEL_ENDPOINT,
-            headers=SERVICE_HEADERS,
-            timeout=15,
-        )
-    else:
-        print("WARNING: no auth settings for Opentelemetry", file=sys.stderr)
-        return
-    return span_exporter
-
-
-def set_zipkin_exporter():
+        config["headers"] = SERVICE_HEADERS
+
+    return OTLPSpanExporter(**config)
+
+
+def _create_zipkin_exporter():
     from opentelemetry.exporter.zipkin.proto.http import ZipkinExporter
 
-    span_exporter = ZipkinExporter(
-        endpoint=ZIPKIN_ENDPOINT,
-    )
-    return span_exporter
+    return ZipkinExporter(endpoint=ZIPKIN_ENDPOINT)
 
 
-def set_console_exporter():
+def _create_console_exporter():
     from opentelemetry.sdk.trace.export import ConsoleSpanExporter
 
-    span_exporter = ConsoleSpanExporter()
-    return span_exporter
+    return ConsoleSpanExporter()

metaflow/tracing/tracing_modules.py

@@ -13,29 +13,31 @@ from typing import Dict, List, Optional
 from opentelemetry import trace as trace_api, context
 from .span_exporter import get_span_exporter
 
-tracer_provider = None
+tracer_provider: Optional[TracerProvider] = None
 
 
 def init_tracing():
     global tracer_provider
+    # Disable logging from opentelemetry
+    import logging
+    logging.getLogger("opentelemetry").setLevel(logging.FATAL)
     if tracer_provider is not None:
-        print("Tracing already initialized", file=sys.stderr)
         return
 
     from .propagator import EnvPropagator
 
     set_global_textmap(EnvPropagator(None))
-    span_exporter = get_span_exporter()
 
-    if "METAFLOW_KUBERNETES_POD_NAMESPACE" in os.environ:
-        service_name = "metaflow-kubernetes"
-    elif "AWS_BATCH_JOB_ID" in os.environ:
-        service_name = "metaflow-awsbatch"
-    else:
-        service_name = "metaflow-local"
+    span_exporter = get_span_exporter()
+    if span_exporter is None:
+        return
 
     tracer_provider = TracerProvider(
-        resource=Resource.create({SERVICE_NAME: service_name})
+        resource=Resource.create(
+            {
+                SERVICE_NAME: "metaflow",
+            }
+        )
     )
     trace_api.set_tracer_provider(tracer_provider)
 
@@ -45,7 +47,9 @@ def init_tracing():
     try:
         from opentelemetry.instrumentation.requests import RequestsInstrumentor
 
-        RequestsInstrumentor().instrument()
+        RequestsInstrumentor().instrument(
+            tracer_provider=tracer_provider,
+        )
     except ImportError:
         # Its possible that this environment has opentelemetry but not requests
         # instrumentation. We don't want to fail in that case.
@@ -55,8 +59,10 @@ def init_tracing():
 @contextlib.contextmanager
 def post_fork():
     global tracer_provider
+
     tracer_provider = None
     init_tracing()
+
     token = context.attach(extract(os.environ))
     try:
         tracer = trace_api.get_tracer_provider().get_tracer(__name__)
@@ -64,47 +70,27 @@ def post_fork():
             "fork", kind=trace_api.SpanKind.SERVER
         ) as span:
             span.set_attribute("cmd", " ".join(sys.argv))
+            span.set_attribute("pid", str(os.getpid()))
             yield
     finally:
         context.detach(token)
 
 
-def _extract_token_after(tokens: List[str], before_token: str) -> Optional[str]:
-    for i, tok in enumerate(tokens):
-        if i > 0 and tokens[i - 1] == before_token:
-            return tok
-
-
-def cli_entrypoint(name: str):
-    def cli_entrypoint_wrap(func):
+def cli(name: str):
+    def cli_wrap(func):
         @wraps(func)
         def wrapper_func(*args, **kwargs):
             global tracer_provider
-
             init_tracing()
 
-            assert tracer_provider is not None  # make type checker happy
+            if tracer_provider is None:
+                return func(*args, **kwargs)
 
             token = context.attach(extract(os.environ))
             try:
                 tracer = trace_api.get_tracer_provider().get_tracer(__name__)
-
-                card_subcommand = _extract_token_after(sys.argv, "card")
-
-                step_name = _extract_token_after(sys.argv, "step")
-                task_id = _extract_token_after(sys.argv, "--task-id")
-                run_id = _extract_token_after(sys.argv, "--run-id")
-                if step_name and task_id and run_id:
-                    better_name = "/".join([run_id, step_name, task_id])
-                elif card_subcommand:
-                    better_name = "card/" + card_subcommand
-                elif "run" in sys.argv:
-                    better_name = "run"
-                else:
-                    better_name = None
-
                 with tracer.start_as_current_span(
-                    better_name or name, kind=trace_api.SpanKind.SERVER
+                    name, kind=trace_api.SpanKind.SERVER
                 ) as span:
                     span.set_attribute("cmd", " ".join(sys.argv))
                     span.set_attribute("pid", str(os.getpid()))
@@ -118,7 +104,7 @@ def cli_entrypoint(name: str):
 
         return wrapper_func
 
-    return cli_entrypoint_wrap
+    return cli_wrap
 
 
 def inject_tracing_vars(env_dict: Dict[str, str]) -> Dict[str, str]:
@@ -127,23 +113,32 @@ def inject_tracing_vars(env_dict: Dict[str, str]) -> Dict[str, str]:
 
 
 def get_trace_id() -> str:
-    return format_trace_id(trace_api.get_current_span().get_span_context().trace_id)
+    try:
+        return format_trace_id(trace_api.get_current_span().get_span_context().trace_id)
+    except Exception:
+        return ""
 
 
 @contextlib.contextmanager
-def traced(name, attrs={}):
+def traced(name: str, attrs: Optional[Dict] = None):
+    if tracer_provider is None:
+        yield
+        return
     tracer = trace_api.get_tracer_provider().get_tracer(__name__)
     with tracer.start_as_current_span(name) as span:
-        for k, v in attrs.items():
-            span.set_attribute(k, v)
+        if attrs:
+            for k, v in attrs.items():
+                span.set_attribute(k, v)
         yield
 
 
 def tracing(func):
     @wraps(func)
     def wrapper_func(*args, **kwargs):
-        tracer = trace_api.get_tracer_provider().get_tracer(func.__module__.__name__)
+        if tracer_provider is None:
+            return func(*args, **kwargs)
 
+        tracer = trace_api.get_tracer_provider().get_tracer(func.__module__)
         with tracer.start_as_current_span(func.__name__):
             return func(*args, **kwargs)
 

metaflow/tuple_util.py

@@ -0,0 +1,27 @@
+# Keep this file minimum dependency as this will be imported by metaflow at bootup.
+def namedtuple_with_defaults(typename, field_descr, defaults=()):
+    from typing import NamedTuple
+
+    T = NamedTuple(typename, field_descr)
+    T.__new__.__defaults__ = tuple(defaults)
+
+    # Adding the following to ensure the named tuple can be (un)pickled correctly.
+    import __main__
+
+    setattr(__main__, T.__name__, T)
+    T.__module__ = "__main__"
+    return T
+
+
+# Define the namedtuple with default here if they need to be accessible in client
+# (and w/o a real flow).
+foreach_frame_field_list = [
+    ("step", str),
+    ("var", str),
+    ("num_splits", int),
+    ("index", int),
+    ("value", str),
+]
+ForeachFrame = namedtuple_with_defaults(
+    "ForeachFrame", foreach_frame_field_list, (None,) * (len(foreach_frame_field_list))
+)