ob-metaflow-extensions 1.1.95__tar.gz → 1.1.97__tar.gz

{ob-metaflow-extensions-1.1.95 → ob-metaflow-extensions-1.1.97}/PKG-INFO

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

{ob-metaflow-extensions-1.1.95 → ob-metaflow-extensions-1.1.97}/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 = [

ob-metaflow-extensions-1.1.97/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.95 → ob-metaflow-extensions-1.1.97}/ob_metaflow_extensions.egg-info/PKG-INFO

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

{ob-metaflow-extensions-1.1.95 → ob-metaflow-extensions-1.1.97}/ob_metaflow_extensions.egg-info/SOURCES.txt

@@ -28,6 +28,7 @@ metaflow_extensions/outerbounds/plugins/snowpark/snowpark_decorator.py
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_exceptions.py
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_service_spec.py
+metaflow_extensions/outerbounds/plugins/tensorboard/__init__.py
 metaflow_extensions/outerbounds/profilers/__init__.py
 metaflow_extensions/outerbounds/profilers/gpu.py
 metaflow_extensions/outerbounds/toplevel/__init__.py

ob-metaflow-extensions-1.1.97/ob_metaflow_extensions.egg-info/requires.txt

@@ -0,0 +1,3 @@
+boto3
+kubernetes
+ob-metaflow==2.12.24.1

{ob-metaflow-extensions-1.1.95 → ob-metaflow-extensions-1.1.97}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_namespace_packages
 from pathlib import Path
 
 
-version = "1.1.95"
+version = "1.1.97"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()
 
@@ -18,6 +18,6 @@ setup(
     install_requires=[
         "boto3",
         "kubernetes",
-        "ob-metaflow == 2.12.23.1",
+        "ob-metaflow == 2.12.24.1",
     ],
 )

ob-metaflow-extensions-1.1.95/ob_metaflow_extensions.egg-info/requires.txt

@@ -1,3 +0,0 @@
-boto3
-kubernetes
-ob-metaflow==2.12.23.1