ob-metaflow-extensions 1.1.96__py2.py3-none-any.whl → 1.1.97__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/__init__.py

@@ -317,6 +317,7 @@ STEP_DECORATORS_DESC = [
         ".fast_bakery.fast_bakery_decorator.InternalFastBakeryDecorator",
     ),
     ("snowpark", ".snowpark.snowpark_decorator.SnowparkDecorator"),
+    ("tensorboard", ".tensorboard.TensorboardDecorator"),
 ]
 FLOW_DECORATORS_DESC = [("nim", ".nim.NimDecorator")]
 TOGGLE_STEP_DECORATOR = [

metaflow_extensions/outerbounds/plugins/tensorboard/__init__.py

@@ -0,0 +1,50 @@
+import functools
+from metaflow.decorators import StepDecorator
+
+
+class TensorboardDecorator(StepDecorator):
+    name = "tensorboard"
+    defaults = {}
+
+    def task_decorate(
+        self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
+    ):
+        @functools.wraps(step_func)
+        def tb_wrapper():
+            import sys, os
+            from metaflow import metaflow_config, current
+
+            try:
+                from torch.utils.tensorboard import SummaryWriter
+            except:
+                print(
+                    "[@tensorboard] Torch and tensorboard not found - logging disabled!",
+                    file=sys.stderr,
+                )
+                step_func()
+            else:
+                tb_root = os.path.join(metaflow_config.DATATOOLS_S3ROOT, "tb")
+                pathspec = current.pathspec
+                try:
+                    log_dir = os.path.join(tb_root, current.project_flow_name, pathspec)
+                except:
+                    log_dir = os.path.join(tb_root, pathspec)
+                comps = log_dir[len(tb_root) + 1 :].split("/")
+                run_level = "/".join(comps[:-2])
+                flow_level = "/".join(comps[:-3])
+
+                print("[@tensorboard] -- INSPECTING RESULTS")
+                print(
+                    "[@tensorboard] -- Execute one of these commands on your workstation:"
+                )
+                print(f"[@tensorboard] Compare tasks of this run: obtb {run_level}")
+                print(f"[@tensorboard] Compare across runs: obtb {flow_level}")
+                writer = SummaryWriter(log_dir=log_dir)
+                setattr(flow, "obtb", writer)
+                try:
+                    step_func()
+                finally:
+                    writer.flush()
+                    delattr(flow, "obtb")
+
+        return tb_wrapper

{ob_metaflow_extensions-1.1.96.dist-info → ob_metaflow_extensions-1.1.97.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.96
+Version: 1.1.97
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob_metaflow_extensions-1.1.96.dist-info → ob_metaflow_extensions-1.1.97.dist-info}/RECORD

@@ -1,7 +1,7 @@
 metaflow_extensions/outerbounds/__init__.py,sha256=TRGvIUMjkfneWtYUFSWoubu_Kf2ekAL4WLbV3IxOj9k,499
 metaflow_extensions/outerbounds/remote_config.py,sha256=Zpfpjgz68_ZgxlXezjzlsDLo4840rkWuZgwDB_5H57U,4059
 metaflow_extensions/outerbounds/config/__init__.py,sha256=JsQGRuGFz28fQWjUvxUgR8EKBLGRdLUIk_buPLJplJY,1225
-metaflow_extensions/outerbounds/plugins/__init__.py,sha256=WwvFcN5kserbPwhrE4hXprnXzJxonEPT0Mlik2kmGMA,12406
+metaflow_extensions/outerbounds/plugins/__init__.py,sha256=7-n1jwyvdNkMo3oU9CMF-YABr6jWRqlI-g8u9LexVR0,12464
 metaflow_extensions/outerbounds/plugins/auth_server.py,sha256=1v2GBqoMBxp5E7Lejz139w-jxJtPnLDvvHXP0HhEIHI,2361
 metaflow_extensions/outerbounds/plugins/perimeters.py,sha256=QXh3SFP7GQbS-RAIxUOPbhPzQ7KDFVxZkTdKqFKgXjI,2697
 metaflow_extensions/outerbounds/plugins/fast_bakery/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -26,6 +26,7 @@ metaflow_extensions/outerbounds/plugins/snowpark/snowpark_decorator.py,sha256=3A
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_exceptions.py,sha256=FTfYlJu-sn9DkPOs2R1V1ChWb1vZthOgeq0BZdT1ucY,296
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py,sha256=d_5UhXqZ_12rCvatH1capPQZYGLx1FVqq_rtW65OXyk,6874
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_service_spec.py,sha256=AI_kcm1hZV3JRxJkookcH6twiGnAYjk9Dx-MeoYz60Y,8511
+metaflow_extensions/outerbounds/plugins/tensorboard/__init__.py,sha256=9lUM4Cqi5RjrHBRfG6AQMRz8-R96eZC8Ih0KD2lv22Y,1858
 metaflow_extensions/outerbounds/profilers/__init__.py,sha256=wa_jhnCBr82TBxoS0e8b6_6sLyZX0fdHicuGJZNTqKw,29
 metaflow_extensions/outerbounds/profilers/gpu.py,sha256=a5YZAepujuP0uDqG9UpXBlZS3wjUt4Yv8CjybXqeT2c,24342
 metaflow_extensions/outerbounds/toplevel/__init__.py,sha256=qWUJSv_r5hXJ7jV_On4nEasKIfUCm6_UjkjXWA_A1Ts,90
@@ -33,7 +34,7 @@ metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,
 metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py,sha256=WUuhz2YQfI4fz7nIcipwwWq781eaoHEk7n4GAn1npDg,63
 metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py,sha256=BbZiaH3uILlEZ6ntBLKeNyqn3If8nIXZFq_Apd7Dhco,70
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
-ob_metaflow_extensions-1.1.96.dist-info/METADATA,sha256=TaJnt4v7k41yQoFx4x2o5lVe1YYO5Z6xFWAKMOtsJw4,520
-ob_metaflow_extensions-1.1.96.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.96.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.96.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.97.dist-info/METADATA,sha256=DuTSqhJsXs0n1kfud4ayP92V0uos7h7ZYsNFc8eZidQ,520
+ob_metaflow_extensions-1.1.97.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.97.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.97.dist-info/RECORD,,