Make g3pdb lazy again

examples/vit_b_i1k.py

@@ -28,7 +28,6 @@
 # pylint: disable=g-import-not-at-top
 with konfig.imports():
   from kauldron import kd
-  from kauldron.data import extra_image_ops as kd_extra
   import optax
 # pylint: enable=g-import-not-at-top
 
@@ -113,7 +112,9 @@ def _make_ds(training: bool):
         kd.data.Elements(keep=["image", "label"]),
         kd.data.InceptionCrop(key="image", resize_size=(224, 224)),
         kd.data.RandomFlipLeftRight(key="image"),
-        kd_extra.RandAugment(image_key="image", num_layers=2, magnitude=15),
+        kd.contrib.data.RandAugment(
+            image_key="image", num_layers=2, magnitude=15
+        ),
         kd.data.ValueRange(key="image", in_vrange=(0, 255), vrange=(0, 1)),
         kd.data.Rearrange(key="label", pattern="... -> ... 1"),
     ]

kauldron/main.py

@@ -19,6 +19,7 @@
 from __future__ import annotations
 
 import contextlib
+import sys
 
 from absl import app
 from absl import flags

kauldron/metrics/classification.py

@@ -19,6 +19,7 @@
 import dataclasses
 from typing import Optional
 
+from etils import epy
 import flax.linen as nn
 import flax.struct
 import jax.numpy as jnp
@@ -27,7 +28,10 @@
 from kauldron.metrics import base_state
 from kauldron.typing import Bool, Float, Int, check_type, typechecked  # pylint: disable=g-multiple-import,g-importing-member
 import numpy as np
-import sklearn.metrics
+
+
+with epy.lazy_imports():
+  import sklearn.metrics  # pylint: disable=g-import-not-at-top
 
 
 @dataclasses.dataclass(kw_only=True, frozen=True, eq=True)

kauldron/summaries/base.py

@@ -22,6 +22,7 @@
 from typing import Any, Mapping, Optional
 
 import einops
+from etils import epy
 import flax
 import jax
 import jax.numpy as jnp
@@ -31,9 +32,11 @@
 import matplotlib
 import mediapy as media
 import numpy as np
-import sklearn.decomposition
 import tensorflow as tf
 
+with epy.lazy_imports():
+  import sklearn.decomposition  # pylint: disable=g-import-not-at-top
+
 
 Images = Float["*b h w c"] | UInt8["*b h w c"]
 Masks = Bool["*b h w 1"]

kauldron/utils/import_test.py

@@ -0,0 +1,33 @@
+# Copyright 2024 The kauldron Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Test Kauldron imports."""
+
+import sys
+import time
+
+
+def test_import_kauldron():
+  start = time.time()
+  from kauldron import kd  # pylint: disable=g-import-not-at-top
+
+  end = time.time()
+
+  # Enforce lazy deps
+  assert "sklearn" not in sys.modules
+
+  # TODO(epot): Reduce this value
+  assert end - start < 40, "Kauldron import took too long."
+
+  del kd

kauldron/xm/_src/kauldron_utils.py

@@ -235,6 +235,7 @@ def base_job(self) -> job_lib.Job:
             "cfg": dir_utils.file_path("config.py"),
             "cfg.workdir": dir_utils.WU_DIR_PROXY,
             **{f"cfg.{k}": v for k, v in self.overrides.items()},
+            "undefok": "g3pdb_port",
         },
         files={
             "config.py": f"//{self.config_path}",