copybara strip some externally unsupported features of the metric writer

PiperOrigin-RevId: 688554082

examples/mnist_autoencoder.py

@@ -14,13 +14,10 @@
 
 r"""Minimal example training a simple Autoencoder on MNIST.
 
-`--xp.use_interpreter` to launch with `ml_python` (no BUILD rules).
-
+Run:
 ```sh
-xmanager launch third_party/py/kauldron/xm/launch.py -- \
-  --cfg=third_party/py/kauldron/examples/mnist_autoencoder.py \
-  --xp.use_interpreter \
-  --xp.platform=jf=2x2
+python main.py --cfg=examples/mnist_autoencoder.py \
+    --cfg.workdir=/tmp/kauldron_oss/workdir
 ```
 
 """
@@ -90,7 +87,8 @@ def get_config():
 
 
 def _make_ds(training: bool):
-  return kd.data.Tfds(
+  Tfds = kd.data.py.Tfds
+  return Tfds(
       name="mnist",
       split="train" if training else "test",
       shuffle=True if training else False,

kauldron/train/metric_writer.py

@@ -358,15 +358,9 @@ class KDMetricWriter(MetadataWriter):
 
   add_artifacts: bool = True
 
-  flatboard_build_context: kdash.BuildContext = (
-      config_util.ROOT_CFG_REF.setup.flatboard_build_context
-  )
-
   @functools.cached_property
   def _collection_path_prefix(self) -> str:
-    if self.flatboard_build_context.collection_path_prefix is None:
-      raise ValueError("collection_path_prefix must be set.")
-    return self.flatboard_build_context.collection_path_prefix
+    return ""
 
   @functools.cached_property
   def _scalar_datatable_name(self) -> str:
@@ -416,27 +410,7 @@ def _array_writer(self) -> metric_writers.MetricWriter:
   def _create_datatable_writer(
       self, name: str, description: str
   ) -> metric_writers.MetricWriter:
-    if epy.is_test():
-      return self._noop  # Do not write to datatable inside tests
-    if status.on_xmanager:
-      if not status.is_lead_host:
-        return self._noop
-      if status.wid == 1 and self.add_artifacts:
-        status.xp.create_artifact(
-            artifact_type=xmanager_api.ArtifactType.ARTIFACT_TYPE_STORAGE2_BIGTABLE,
-            artifact=name,
-            description=description,
-        )
-      keys = [("wid", status.wid)]
-    else:
-      keys = []
-
-    return metric_writers.AsyncWriter(
-        metric_writers.DatatableWriter(
-            datatable_name=name,
-            keys=keys,
-        ),
-    )
+    return self._noop
 
   def write_summaries(
       self,
@@ -504,12 +478,9 @@ def write_pointcloud(
       point_colors: Mapping[str, Array["n 3"]] | None = None,
       configs: Mapping[str, str | float | bool | None] | None = None,
   ) -> None:
-    self._tf_summary_writer.write_pointcloud(
-        step=step,
-        point_clouds=point_clouds,
-        point_colors=point_colors,
-        configs=configs,
-    )
+    if not point_clouds:
+      return
+    logging.info("Pointcloud summary not supported.")
 
   def write_hparams(self, hparams: Mapping[str, Any]) -> None:
     self._log_writer.write_hparams(hparams)

pyproject.toml

@@ -42,8 +42,10 @@ dependencies = [
     "regex",
     "scikit-image",
     "scikit-learn",
+    "tabulate",  # used by pandas.DataFrame.to_markdown (for logging context)
     "tensorflow",
-    "tensorflow_datasets",
+    "tfds-nightly",  # TODO(klausg): switch back to tensorflow_datasets>=4.9.7
+    # once released: https://github.com/tensorflow/datasets/commit/d4bfd59863c6cb5b64d043b7cb6ab566e7d92440
     "tqdm",
     # closest match to the internal typeguard
     "typeguard@git+https://github.com/agronholm/typeguard@0dd7f7510b7c694e66a0d17d1d58d185125bad5d",