From 541b4b9f8806ceaf894ecb9b628d19574cfe03ab Mon Sep 17 00:00:00 2001 From: Drew Date: Wed, 28 Oct 2026 12:07:14 -0700 Subject: [PATCH 1/3] Remove the data_sample model constructor argument Hyrax now constructs every model as Model(config) and passes nothing else. The pre-flight layer that existed only to feed a data_sample is deleted: DataProvider.sample_data, StreamingDataProvider.sample_data, HyraxCSVDataset.sample_data, the Kafka and LSDB peek_sample methods with their replay buffers, the peek_timeout config key, setup_model_from_sample, and the reset_trace call that cleaned up after the pre-flight. The built-in image models (HyraxCNN, HyraxAutoencoder, HyraxAutoencoderV2, ImageDCAE) size their layers from [model.].input_shape, a [channels, height, width] list defaulting to [3, 32, 32]. A shared read_input_shape helper validates it and raises a ValueError naming the key. Height and width are now assigned correctly where they were swapped. HyraxTs2Vec sizes its input from EncodingSettings, and SimCLR is constructable through setup_model again. infer_stream no longer takes sample_batch: a configured [data_request.infer_stream] group gives an iterable session, otherwise a manual session fed through process(). Migration 007 drops peek_timeout from existing configs. Docs, the model reference, and the maintained notebooks are updated to the new contract. Co-Authored-By: Claude Fable 5.1 --- HYRAX_GUIDE.md | 2 + docs/model_class_reference.rst | 31 +++---- docs/notebooks/custom_training_metrics.ipynb | 4 +- .../list_available_models_and_datasets.ipynb | 4 +- docs/pre_executed/external_model_class.ipynb | 27 +++--- docs/pre_executed/lsdb_stream_dataset.ipynb | 17 ++-- .../supervised_lightcurve_transients.ipynb | 6 +- .../007_remove_kafka_peek_timeout.py | 24 ++++++ src/hyrax/context.py | 2 +- src/hyrax/datasets/data_provider.py | 11 --- src/hyrax/datasets/hyrax_csv_dataset.py | 9 -- src/hyrax/datasets/kafka_stream_dataset.py | 86 +------------------ src/hyrax/datasets/lsdb_stream_dataset.py | 48 ++--------- src/hyrax/datasets/streaming_data_provider.py | 16 +--- src/hyrax/hyrax_default_config.toml | 29 +++++-- src/hyrax/models/hsc_autoencoder.py | 2 +- src/hyrax/models/hsc_dcae.py | 2 +- src/hyrax/models/hyrax_autoencoder.py | 18 ++-- src/hyrax/models/hyrax_autoencoderv2.py | 18 ++-- src/hyrax/models/hyrax_cnn.py | 16 ++-- src/hyrax/models/hyrax_loopback.py | 2 +- src/hyrax/models/hyrax_ts2vec.py | 29 ++----- src/hyrax/models/image_dcae.py | 20 ++--- src/hyrax/models/model_utils.py | 41 +++++++++ src/hyrax/models/simclr.py | 3 +- src/hyrax/pytorch_ignite.py | 58 +------------ src/hyrax/trace.py | 16 +--- src/hyrax/verbs/infer.py | 2 +- src/hyrax/verbs/infer_stream.py | 54 +++--------- src/hyrax/verbs/test.py | 2 +- src/hyrax/verbs/to_onnx.py | 6 +- src/hyrax/verbs/train.py | 2 +- src/hyrax/verbs/train_stream.py | 10 +-- tests/hyrax/conftest.py | 2 +- .../test_migration_007.py | 34 ++++++++ tests/hyrax/test_csv_dataset.py | 23 ----- tests/hyrax/test_data_provider.py | 80 ----------------- tests/hyrax/test_e2e.py | 8 +- tests/hyrax/test_hyrax_autoencoderv2.py | 16 ++-- tests/hyrax/test_hyrax_cnn.py | 17 ++-- tests/hyrax/test_hyrax_ts2vec.py | 13 +-- tests/hyrax/test_infer_stream.py | 23 +++-- tests/hyrax/test_input_shape.py | 77 +++++++++++++++++ tests/hyrax/test_kafka_stream_dataset.py | 77 ----------------- tests/hyrax/test_list_registries.py | 2 +- tests/hyrax/test_lsdb_stream_dataset.py | 32 ------- tests/hyrax/test_model_registry.py | 24 +++--- tests/hyrax/test_patch_model.py | 12 +-- tests/hyrax/test_plugin_utils.py | 8 +- tests/hyrax/test_prepare_inputs_migration.py | 12 +-- tests/hyrax/test_streaming_data_provider.py | 17 +--- tests/hyrax/test_train_stream.py | 2 +- 52 files changed, 401 insertions(+), 695 deletions(-) create mode 100644 src/hyrax/config_migrations/migrations/007_remove_kafka_peek_timeout.py create mode 100644 tests/hyrax/test_config_migrations/test_migration_007.py create mode 100644 tests/hyrax/test_input_shape.py diff --git a/HYRAX_GUIDE.md b/HYRAX_GUIDE.md index 19e22833b..2ca327445 100644 --- a/HYRAX_GUIDE.md +++ b/HYRAX_GUIDE.md @@ -162,6 +162,8 @@ Hyrax discovers components through three registries: - Decorator: `@hyrax_model` - Models must inherit from `torch.nn.Module` and implement `__init__`, `forward`, `train_batch`, and `prepare_inputs`. +- Hyrax constructs a model as `Model(config)` and passes nothing else. Models size + themselves from config; the built-in image models read `[model.].input_shape`. - The decorator wires up save/load, optimizer, and criterion handling. - Every value in the metric dict returned by `train_batch`, `validate_batch`, or `test_batch` must be a per-sample mean over the batch (not a sum or count); epoch metrics weight each diff --git a/docs/model_class_reference.rst b/docs/model_class_reference.rst index 98188f589..8a03580ae 100644 --- a/docs/model_class_reference.rst +++ b/docs/model_class_reference.rst @@ -12,10 +12,8 @@ How Hyrax uses your model class Hyrax does this sequence: -1. Builds one sample batch from your dataset. -2. Calls ``YourModel.prepare_inputs(sample_batch)``. -3. Creates your model with ``YourModel(config, data_sample=prepared_sample)``. -4. During runs, calls: +1. Creates your model with ``YourModel(config)``. +2. During runs, calls: * ``train_batch`` for training * ``infer_batch`` for inference @@ -33,7 +31,7 @@ Your class must have these: find it by name in the :doc:`configuration `). 3. Implement: - * ``__init__(self, config, data_sample=None)`` + * ``__init__(self, config)`` * ``forward(self, ...)`` (required by PyTorch modules) * ``infer_batch(self, batch)`` * ``train_batch(self, batch)`` @@ -44,15 +42,15 @@ Your class must have these: Method-by-method guide ---------------------- -``__init__(self, config, data_sample=None)`` -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +``__init__(self, config)`` +^^^^^^^^^^^^^^^^^^^^^^^^^^ What to do in this method: 1. Save ``config`` to ``self.config``. -2. Read model-specific settings from config. +2. Read model-specific settings from config, including any input sizes your + layers need. Hyrax passes nothing but ``config``, so shapes come from it. 3. Build your layers. -4. Optionally use ``data_sample`` to infer shapes. The model-specific config values are typically defined in your external library config layout. See :doc:`/external_library_package` for packaging/config setup. @@ -61,18 +59,14 @@ Example (read settings from config): .. code-block:: python - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config + in_dim = config["my_package"]["MyModel"]["input_size"] hidden = config["my_package"]["MyModel"]["hidden_size"] out_dim = config["my_package"]["MyModel"]["num_classes"] - in_dim = 128 - if data_sample is not None: - x, _ = data_sample - in_dim = x.shape[-1] - self.fc1 = nn.Linear(in_dim, hidden) self.fc2 = nn.Linear(hidden, out_dim) @@ -80,7 +74,8 @@ To set those config values in a notebook, one simple pattern is: .. code-block:: python - h.config.setdefault("my_package", {}).setdefault("MyModel", {})["hidden_size"] = 64 + h.config["my_package"]["MyModel"]["input_size"] = 128 + h.config["my_package"]["MyModel"]["hidden_size"] = 64 h.config["my_package"]["MyModel"]["num_classes"] = 10 PyTorch reference for building ``nn.Module`` layers: @@ -366,7 +361,7 @@ when you use it always gives you the run that is actually underway. @hyrax_model class MyModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.activations = [] @@ -409,7 +404,7 @@ Complete minimal class @hyrax_model class MyModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.fc = nn.Linear(128, 10) diff --git a/docs/notebooks/custom_training_metrics.ipynb b/docs/notebooks/custom_training_metrics.ipynb index b948639a6..996fe472c 100644 --- a/docs/notebooks/custom_training_metrics.ipynb +++ b/docs/notebooks/custom_training_metrics.ipynb @@ -57,7 +57,7 @@ "class MetricsDemoAutoencoder(nn.Module):\n", " \"\"\"Tiny autoencoder that demonstrates custom training metrics.\"\"\"\n", "\n", - " def __init__(self, config, data_sample=None):\n", + " def __init__(self, config):\n", " super().__init__()\n", " self.encoder = nn.Linear(4, 2)\n", " self.decoder = nn.Linear(2, 4)\n", @@ -135,7 +135,7 @@ "class EpochMetricsDemoAutoencoder(nn.Module):\n", " \"\"\"Autoencoder that reports both batch and epoch-level metrics.\"\"\"\n", "\n", - " def __init__(self, config, data_sample=None):\n", + " def __init__(self, config):\n", " super().__init__()\n", " self.encoder = nn.Linear(4, 2)\n", " self.decoder = nn.Linear(2, 4)\n", diff --git a/docs/notebooks/list_available_models_and_datasets.ipynb b/docs/notebooks/list_available_models_and_datasets.ipynb index 21148edc6..5a9a7d778 100644 --- a/docs/notebooks/list_available_models_and_datasets.ipynb +++ b/docs/notebooks/list_available_models_and_datasets.ipynb @@ -95,7 +95,7 @@ "class AAA_MyCustomModel(nn.Module):\n", " \"\"\"A simple toy model for demonstration purposes.\"\"\"\n", "\n", - " def __init__(self, config, data_sample=None):\n", + " def __init__(self, config):\n", " super().__init__()\n", " self.linear = nn.Linear(10, 2)\n", "\n", @@ -166,7 +166,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.12" + "version": "3.13.15" } }, "nbformat": 4, diff --git a/docs/pre_executed/external_model_class.ipynb b/docs/pre_executed/external_model_class.ipynb index b2ea63046..e9735981d 100644 --- a/docs/pre_executed/external_model_class.ipynb +++ b/docs/pre_executed/external_model_class.ipynb @@ -115,7 +115,7 @@ "\n", "The `@hyrax_model` decorator registers the class by name so Hyrax can find it when you set `model.name`.\n", "\n", - "The constructor receives `config` (the Hyrax config dict) and an optional `data_sample` — a single batch produced by `prepare_inputs`. The sample lets you infer input dimensions automatically instead of hardcoding them." + "The constructor receives `config`, the Hyrax config dict. Any sizes your layers need (image shape, number of classes, sequence length) come from a config entry you add for your model, as shown below with `h.config[\"model\"][\"ImageOnlyCNN\"]`." ] }, { @@ -132,12 +132,11 @@ "\n", "@hyrax_model # registers the class so h.set_config(\"model.name\", \"ImageOnlyCNN\") works\n", "class ImageOnlyCNN(nn.Module):\n", - " def __init__(self, config, data_sample=None):\n", + " def __init__(self, config):\n", " super().__init__()\n", " self.config = config\n", - " # Unpack the sample to derive input dimensions automatically\n", - " image_sample, _ = data_sample\n", - " c, h, w = image_sample.shape[1:] # channels, height, width\n", + " # Input dimensions come from the config entry set in the training cell below\n", + " c, h, w = config[\"model\"][\"ImageOnlyCNN\"][\"input_shape\"] # channels, height, width\n", " self.net = nn.Sequential(\n", " nn.Conv2d(c, 8, kernel_size=3, padding=1),\n", " nn.ReLU(),\n", @@ -201,6 +200,8 @@ "\n", "h = Hyrax()\n", "h.set_config(\"model.name\", \"ImageOnlyCNN\")\n", + "# The image shape the model is built for; the dataset above produces 3x32x32 images.\n", + "h.set_config(\"model.ImageOnlyCNN.input_shape\", [3, 32, 32])\n", "h.set_config(\"train.epochs\", 1)\n", "# These are set to default values, uncomment to change.\n", "# h.set_config(\"data_loader.batch_size\", 512)\n", @@ -255,14 +256,13 @@ "source": [ "@hyrax_model\n", "class ImageAndLightCurveModel(nn.Module):\n", - " def __init__(self, config, data_sample=None):\n", + " def __init__(self, config):\n", " super().__init__()\n", " self.config = config\n", - " image, light_curve, light_curve_mask, _ = data_sample\n", + " model_config = config[\"model\"][\"ImageAndLightCurveModel\"]\n", "\n", - " c, h, w = image.shape[1:]\n", - " lc_dim = light_curve.shape[1] # padded sequence length\n", - " print(f\"LC_DIM={lc_dim}\")\n", + " c, h, w = model_config[\"input_shape\"]\n", + " lc_dim = model_config[\"light_curve_length\"] # padded sequence length\n", "\n", " # Image branch: simple CNN → 32-d embedding\n", " self.image_encoder = nn.Sequential(\n", @@ -332,6 +332,9 @@ "source": [ "h2 = Hyrax()\n", "h2.set_config(\"model.name\", \"ImageAndLightCurveModel\")\n", + "# The sizes the model is built for: 3x32x32 images, and light curves padded to max_len (100).\n", + "h2.set_config(\"model.ImageAndLightCurveModel.input_shape\", [3, 32, 32])\n", + "h2.set_config(\"model.ImageAndLightCurveModel.light_curve_length\", 100)\n", "h2.set_config(\"train.epochs\", 1)\n", "# These are set to default values, uncomment to change.\n", "# h.set_config(\"data_loader.batch_size\", 512)\n", @@ -403,7 +406,7 @@ ], "metadata": { "kernelspec": { - "display_name": "kbmod", + "display_name": "hyrax", "language": "python", "name": "python3" }, @@ -417,7 +420,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.11" + "version": "3.13.15" }, "nbsphinx": { "execute": "never" diff --git a/docs/pre_executed/lsdb_stream_dataset.ipynb b/docs/pre_executed/lsdb_stream_dataset.ipynb index aa7d968d2..8b533a697 100644 --- a/docs/pre_executed/lsdb_stream_dataset.ipynb +++ b/docs/pre_executed/lsdb_stream_dataset.ipynb @@ -277,16 +277,19 @@ }, "outputs": [], "source": [ + "# The number of catalog columns prepare_inputs stacks into each input row.\n", + "N_FEATURES = 2\n", + "\n", + "\n", "@hyrax_model\n", "class CatalogAutoencoder(nn.Module):\n", " \"\"\"A tiny dense autoencoder over a handful of catalog columns.\"\"\"\n", "\n", - " def __init__(self, config, data_sample=None):\n", + " def __init__(self, config):\n", " super().__init__()\n", " self.config = config\n", - " n_features = data_sample.shape[1]\n", - " self.encoder = nn.Sequential(nn.Linear(n_features, 8), nn.GELU(), nn.Linear(8, 2))\n", - " self.decoder = nn.Sequential(nn.Linear(2, 8), nn.GELU(), nn.Linear(8, n_features))\n", + " self.encoder = nn.Sequential(nn.Linear(N_FEATURES, 8), nn.GELU(), nn.Linear(8, 2))\n", + " self.decoder = nn.Sequential(nn.Linear(2, 8), nn.GELU(), nn.Linear(8, N_FEATURES))\n", "\n", " def forward(self, batch):\n", " return self.encoder(batch)\n", @@ -304,7 +307,7 @@ "\n", " @staticmethod\n", " def prepare_inputs(data_dict) -> tuple:\n", - " \"\"\"Stack the requested catalog columns into one (batch, n_features) array.\"\"\"\n", + " \"\"\"Stack the requested catalog columns into one (batch, N_FEATURES) array.\"\"\"\n", " import numpy as np\n", "\n", " data = data_dict[\"data\"]\n", @@ -651,7 +654,7 @@ ], "metadata": { "kernelspec": { - "display_name": "hyrax-lsdb-dataset", + "display_name": "hyrax", "language": "python", "name": "python3" }, @@ -665,7 +668,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.14" + "version": "3.13.15" } }, "nbformat": 4, diff --git a/docs/pre_executed/supervised_lightcurve_transients.ipynb b/docs/pre_executed/supervised_lightcurve_transients.ipynb index cb76db2ba..f44b7b206 100644 --- a/docs/pre_executed/supervised_lightcurve_transients.ipynb +++ b/docs/pre_executed/supervised_lightcurve_transients.ipynb @@ -514,7 +514,7 @@ "class LightCurveGRU(nn.Module):\n", " \"\"\"A bidirectional GRU for variable-length multi-band light-curve classification.\"\"\"\n", "\n", - " def __init__(self, config, data_sample=None):\n", + " def __init__(self, config):\n", " super().__init__()\n", " self.config = config\n", " model_config = config[\"model\"][\"LightCurveGRU\"]\n", @@ -1623,7 +1623,7 @@ ], "metadata": { "kernelspec": { - "display_name": "kbmod", + "display_name": "hyrax", "language": "python", "name": "python3" }, @@ -1637,7 +1637,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.11" + "version": "3.13.15" }, "widgets": { "application/vnd.jupyter.widget-state+json": { diff --git a/src/hyrax/config_migrations/migrations/007_remove_kafka_peek_timeout.py b/src/hyrax/config_migrations/migrations/007_remove_kafka_peek_timeout.py new file mode 100644 index 000000000..eee0e7169 --- /dev/null +++ b/src/hyrax/config_migrations/migrations/007_remove_kafka_peek_timeout.py @@ -0,0 +1,24 @@ +"""Config migration: version 7 → version 8. + +Removes ``[data_set.KafkaStreamDataset].peek_timeout``. Models are no longer +pre-flighted from a sample peeked off the stream, so there is no wait left for +the key to bound. +""" + +from tomlkit.toml_document import TOMLDocument + +from hyrax.config_migrations.migration_utils import migration_step + + +@migration_step(from_version=7) +def remove_kafka_peek_timeout(cfg: TOMLDocument) -> TOMLDocument: + """Drop the obsolete ``peek_timeout`` key from ``[data_set.KafkaStreamDataset]``.""" + data_set = cfg.get("data_set") + if not isinstance(data_set, dict): + return cfg + + kafka = data_set.get("KafkaStreamDataset") + if isinstance(kafka, dict): + kafka.pop("peek_timeout", None) + + return cfg diff --git a/src/hyrax/context.py b/src/hyrax/context.py index 5f3dff301..9f3e8cff7 100644 --- a/src/hyrax/context.py +++ b/src/hyrax/context.py @@ -16,7 +16,7 @@ @hyrax_model class MyModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config diff --git a/src/hyrax/datasets/data_provider.py b/src/hyrax/datasets/data_provider.py index b5feca141..3dbf26561 100644 --- a/src/hyrax/datasets/data_provider.py +++ b/src/hyrax/datasets/data_provider.py @@ -1053,17 +1053,6 @@ def _apply_configurations(base_config: dict, dataset_definition: dict) -> dict: else: return base_config - def sample_data(self) -> dict: - """Returns a data sample. Primarily this will be used for instantiating a - model so that any runtime resizing can be handled properly. - - Returns - ------- - dict - A dictionary containing the data for index 0. - """ - return self[0] - def get_object_id(self, idx: int) -> str: """Returns the ID at a particular index. diff --git a/src/hyrax/datasets/hyrax_csv_dataset.py b/src/hyrax/datasets/hyrax_csv_dataset.py index d34ce5afa..e905e2dde 100644 --- a/src/hyrax/datasets/hyrax_csv_dataset.py +++ b/src/hyrax/datasets/hyrax_csv_dataset.py @@ -71,12 +71,3 @@ def __getitem__(self, idx): def __len__(self) -> int: """Return the number of records in the CSV.""" return len(self.mem_mapped_csv) - - def sample_data(self): - """Return the first record, in dictionary form, as the sample.""" - sample = {"data": {}} - - for col in self.column_names: - sample["data"][col] = self.mem_mapped_csv.iloc[0][col] - - return sample diff --git a/src/hyrax/datasets/kafka_stream_dataset.py b/src/hyrax/datasets/kafka_stream_dataset.py index 42374025f..1bdcaa8f5 100644 --- a/src/hyrax/datasets/kafka_stream_dataset.py +++ b/src/hyrax/datasets/kafka_stream_dataset.py @@ -76,7 +76,6 @@ import os import socket import threading -import time from pathlib import Path from urllib.parse import urlparse @@ -87,9 +86,6 @@ logger = logging.getLogger(__name__) -# How often to log a "still waiting" warning while peek_sample() blocks on an empty topic. -PEEK_WARN_INTERVAL = 15.0 - def _format_partitions(partitions) -> list[str]: """Render a list of ``TopicPartition`` as ``["topic[0]", "topic[1]"]`` for logging.""" @@ -137,21 +133,14 @@ def __init__(self, config: dict, data_location=None): self.batch_flush_timeout = float(ds_config["batch_flush_timeout"]) self.commit_after_batch = bool(ds_config["commit_after_batch"]) - # Bound on how long peek_sample() waits for its first message. The TOML `false` - # sentinel (or 0) means "wait indefinitely". - peek_timeout = ds_config["peek_timeout"] - self.peek_timeout = float(peek_timeout) if peek_timeout else None - # The flush threshold is the configured DataLoader batch size. self.batch_size = config["data_loader"]["batch_size"] # Set from another thread (or session teardown) to end iteration; see stop(). self._stop = threading.Event() - # Single shared consumer, created lazily. Shared between peek_sample() and - # __iter__ so a peeked message can be replayed into the first batch. + # Single consumer, created lazily on first iteration and released by close(). self._consumer = None - self._buffered: list[dict] = [] credentials_file = ds_config.get("credentials_file") credentials_file_path = Path(credentials_file) if credentials_file else None @@ -354,73 +343,6 @@ def _decode(self, msg) -> dict: """ return json.loads(msg.value()) - def peek_sample(self) -> dict: - """Return one decoded sample without removing it from the batch stream. - - Polls until a message arrives (or :meth:`stop` is set), decodes it, and buffers - it so :meth:`__iter__` replays it as part of the first batch. Used to pre-flight - the model architecture without losing a live message. - - Waiting is bounded by ``peek_timeout`` and reports what the consumer is connected - to when it gives up. An unbounded silent wait here is indistinguishable from an - unreachable broker, an empty topic, or a partition assignment lost to a consumer - left over from an earlier run. - - Returns - ------- - dict - The flat decoded sample. - - Raises - ------ - TimeoutError - If ``peek_timeout`` elapses before a message arrives. - RuntimeError - If the stream is stopped before any message arrives. - """ - consumer = self._ensure_consumer() - started = time.monotonic() - deadline = started + self.peek_timeout if self.peek_timeout else None - next_warning = started + PEEK_WARN_INTERVAL - - try: - while not self._stop.is_set(): - now = time.monotonic() - if deadline is not None and now >= deadline: - raise TimeoutError( - self._connection_diagnostic( - f"No Kafka message arrived within peek_timeout ({self.peek_timeout}s)" - ) - ) - if now >= next_warning: - next_warning = now + PEEK_WARN_INTERVAL - logger.warning( - self._connection_diagnostic( - f"Still waiting for a Kafka message after {now - started:.0f}s" - ) - ) - - msg = consumer.poll(self.poll_timeout) - if msg is None: - continue - if msg.error() is not None: - self._log_message_error(msg) - continue - - sample = self._decode(msg) - self._buffered.append(sample) - return sample - - raise RuntimeError("KafkaStreamDataset.peek_sample() stopped before a message arrived.") - except BaseException: - # There is no `finally` here on purpose: on success the consumer must stay - # open for __iter__. On failure it must not -- and BaseException rather than - # Exception because interrupting a hung notebook cell is the common case, and - # a KeyboardInterrupt that leaked the consumer is exactly how a zombie group - # member is created. - self.close() - raise - def __iter__(self): """Poll Kafka and yield ``list[dict]`` batches with latency-bounded flushing. @@ -437,10 +359,6 @@ def __iter__(self): consumer = self._ensure_consumer() - # Replay any peeked-but-not-yet-delivered messages ahead of the next batch. - pending: list[dict] = list(self._buffered) - self._buffered = [] - try: while not self._stop.is_set(): # consume() blocks up to batch_flush_timeout and returns up to batch_size @@ -449,7 +367,7 @@ def __iter__(self): # until the batch fills. messages = consumer.consume(num_messages=self.batch_size, timeout=self.batch_flush_timeout) - batch, pending = pending, [] + batch: list[dict] = [] for msg in messages: if msg.error() is None: batch.append(self._decode(msg)) diff --git a/src/hyrax/datasets/lsdb_stream_dataset.py b/src/hyrax/datasets/lsdb_stream_dataset.py index dc82c2ca5..f141bc09b 100644 --- a/src/hyrax/datasets/lsdb_stream_dataset.py +++ b/src/hyrax/datasets/lsdb_stream_dataset.py @@ -175,12 +175,11 @@ def __init__(self, config: dict, data_location=None): # Set from another thread (or session teardown) to end iteration; see stop(). self._stop = threading.Event() - # Single shared iterator, created lazily. Shared between peek_sample() and - # __iter__ so peeked rows can be replayed into the first batch. + # Single shared iterator, created lazily, so successive passes of __iter__ resume + # the same traversal. Rows fetched but not yet delivered wait in _buffered. self._iterator = None self._exhausted = False self._buffered: list[dict] = [] - self._peek_index = 0 # A dask Client this dataset created (and therefore owns) via # A client attached via get_client() is owned by the caller and is never @@ -408,8 +407,8 @@ def _ensure_iterator(self): """Return the shared stream iterator, creating it on first use. ``CatalogStream.__iter__`` hands back a *new* iterator with a freshly spawned RNG - on every call, so peek_sample() and __iter__() must share a single one or peeked - rows would be drawn from a different traversal than the batches. + on every call, so successive passes of __iter__() must share a single one or the + rows one pass pushed back would come from a different traversal than the next. """ if self._iterator is None: self._exhausted = False @@ -465,42 +464,6 @@ def __len__(self): # Row production # - def peek_sample(self) -> dict: - """Return one row without removing it from the batch stream. - - Pulls chunks until a row is available, buffers them so :meth:`__iter__` replays - them as part of the first batch, and advances a cursor so repeated calls return - distinct rows. Used to pre-flight the model architecture without losing data. - - Returns - ------- - dict - The flat row, mapping column name to value. - - Raises - ------ - RuntimeError - If the stream is stopped, or produces no rows, before one is available. - """ - iterator = self._ensure_iterator() - - # A `while`, not an `if`: a partition can legitimately compute to zero rows. - while self._peek_index >= len(self._buffered): - if self._stop.is_set(): - raise RuntimeError("LSDBStreamDataset.peek_sample() was stopped before a row arrived.") - try: - chunk = next(iterator) - except StopIteration as err: - raise RuntimeError( - f"LSDBStreamDataset.peek_sample(): the catalog at '{self.data_location}' " - "produced no rows." - ) from err - self._buffered.extend(chunk.to_dict(orient="records")) - - sample = self._buffered[self._peek_index] - self._peek_index += 1 - return sample - def __iter__(self): """Yield ``list[dict]`` batches of exactly ``batch_size`` rows. @@ -514,10 +477,9 @@ def __iter__(self): """ iterator = self._ensure_iterator() - # Replay any peeked-but-not-yet-delivered rows ahead of the first batch. + # Resume with the rows an earlier pass fetched but never delivered. batch: list[dict] = list(self._buffered) self._buffered = [] - self._peek_index = 0 # The chunk currently being drained, and how far into it we have read. Tracked out # here so the `finally` can push back the rows of a half-drained chunk. diff --git a/src/hyrax/datasets/streaming_data_provider.py b/src/hyrax/datasets/streaming_data_provider.py index 158aed7ec..0c6ad4e75 100644 --- a/src/hyrax/datasets/streaming_data_provider.py +++ b/src/hyrax/datasets/streaming_data_provider.py @@ -4,10 +4,10 @@ random access (``resolve_data(idx)``, joins, caching, augmentation, splits), none of which a live stream supports. :class:`StreamingDataProvider` is its **iterable** sibling: it wraps a single streaming dataset (e.g. -:class:`~hyrax.datasets.kafka_stream_dataset.KafkaStreamDataset`) and presents the same -surface the rest of Hyrax expects — ``collate`` (shared via -:class:`~hyrax.datasets.data_provider.CollationMixin`) and ``sample_data`` for model -pre-flighting — while delegating iteration to the wrapped stream. +:class:`~hyrax.datasets.kafka_stream_dataset.KafkaStreamDataset`) and presents the +``collate`` surface the rest of Hyrax expects (shared via +:class:`~hyrax.datasets.data_provider.CollationMixin`) while delegating iteration to the +wrapped stream. The wrapped stream is a *dumb decoder* that yields ``list[dict]`` batches of flat sample dicts. This provider owns the structuring: it extracts the object id (via @@ -151,14 +151,6 @@ def __iter__(self): for batch in self._stream: yield [self._structure(sample) for sample in batch] - def sample_data(self) -> dict: - """Return one structured sample for model pre-flighting (``setup_model``). - - Peeks a single message from the stream without losing it (it is replayed in the - first batch) and structures it like any other sample. - """ - return self._structure(self._stream.peek_sample()) - def stop(self): """Stop the underlying stream's iteration.""" self._stream.stop() diff --git a/src/hyrax/hyrax_default_config.toml b/src/hyrax/hyrax_default_config.toml index 9960dd01c..a1eebe866 100644 --- a/src/hyrax/hyrax_default_config.toml +++ b/src/hyrax/hyrax_default_config.toml @@ -90,24 +90,39 @@ name = "" [model.HyraxAutoencoder] +# The shape of one input image as [channels, height, width]. The layers are sized from +# this when the model is built, so it must match what the dataset produces (see +# [data_set].crop_to). The default matches HyraxCifarDataset. +input_shape = [3, 32, 32] + # The number of output channels from the first layer. base_channel_size = 32 -# The length of the latent space vector. +# The length of the latent space vector. latent_dim = 64 [model.HyraxAutoencoderV2] +# The shape of one input image as [channels, height, width]. The layers are sized from +# this when the model is built, so it must match what the dataset produces (see +# [data_set].crop_to). The default matches HyraxCifarDataset. +input_shape = [3, 32, 32] + # The number of output channels from the first layer. base_channel_size = 32 -# The length of the latent space vector. +# The length of the latent space vector. latent_dim = 64 # The activation function of the final layer. final_layer = "tanh" [model.ImageDCAE] +# The shape of one input image as [channels, height, width]. The layers are sized from +# this when the model is built, so it must match what the dataset produces (see +# [data_set].crop_to). The default matches HyraxCifarDataset. +input_shape = [3, 32, 32] + # The number of output channels from the first layer. base_channel_size = 32 @@ -146,6 +161,11 @@ gaussian_blur_sigma_range = [0.1, 2.0] [model.HyraxCNN] +# The shape of one input image as [channels, height, width]. The layers are sized from +# this when the model is built, so it must match what the dataset produces (see +# [data_set].crop_to). The default matches HyraxCifarDataset. +input_shape = [3, 32, 32] + # The number of classes to predict as the output of the model. i.e. 2 would be a # binary classifer, 10 would predict the 10 classes in the CiFAR dataset. output_classes = 10 @@ -467,11 +487,6 @@ poll_timeout = 1.0 # Maximum seconds to wait before flushing a partial (smaller than batch_size) batch. batch_flush_timeout = 5.0 -# Maximum seconds peek_sample() waits for a first message before raising TimeoutError. -# This bounds model pre-flighting, which would otherwise block forever against an empty -# topic or an assignment held by another consumer. Set to false to wait indefinitely. -peek_timeout = 120.0 - # Provided path credentials file for a `confluent_kafka` consumer instance # e.g. # "sasl.username" = diff --git a/src/hyrax/models/hsc_autoencoder.py b/src/hyrax/models/hsc_autoencoder.py index 9842df349..f03ef5643 100644 --- a/src/hyrax/models/hsc_autoencoder.py +++ b/src/hyrax/models/hsc_autoencoder.py @@ -13,7 +13,7 @@ class HSCAutoencoder(nn.Module): # These shapes work with [3,258,258] inputs This autoencoder is designed to work with datasets that are prepared with Hyrax's HSC Data Set class. """ - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() # Encoder diff --git a/src/hyrax/models/hsc_dcae.py b/src/hyrax/models/hsc_dcae.py index eeaeb754a..70f9d5c73 100644 --- a/src/hyrax/models/hsc_dcae.py +++ b/src/hyrax/models/hsc_dcae.py @@ -24,7 +24,7 @@ class HSCDCAE(nn.Module): This autoencoder is designed to work with datasets that are prepared with Hyrax's HSC Data Set class. """ - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() # The current network works with images of size [3,150,150] diff --git a/src/hyrax/models/hyrax_autoencoder.py b/src/hyrax/models/hyrax_autoencoder.py index 30b42e84a..3856750b3 100644 --- a/src/hyrax/models/hyrax_autoencoder.py +++ b/src/hyrax/models/hyrax_autoencoder.py @@ -8,6 +8,7 @@ # extra long import here to address a circular import issue from hyrax.models.model_registry import hyrax_model +from hyrax.models.model_utils import read_input_shape logger = logging.getLogger(__name__) @@ -21,18 +22,17 @@ class HyraxAutoencoder(nn.Module): `autoenocoder tutorial `_ The train function has been converted into train_batch for use with pytorch-ignite. + + The layers are sized from ``[model.HyraxAutoencoder].input_shape``, the + ``[channels, height, width]`` of one input image. """ - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config - shape = data_sample.shape - logger.debug(f"Found shape: {shape} in data sample, using this to initialize model.") - - # Unpack the shape of the image (batch_size, num_channels, width, height) - # we'll ignore the batch_size during initialization. - _, self.num_input_channels, self.image_width, self.image_height = shape + self.input_shape = read_input_shape(config, "HyraxAutoencoder") + self.num_input_channels, self.image_height, self.image_width = self.input_shape self.c_hid = self.config["model"]["HyraxAutoencoder"]["base_channel_size"] self.latent_dim = self.config["model"]["HyraxAutoencoder"]["latent_dim"] @@ -41,8 +41,8 @@ def __init__(self, config, data_sample=None): # Formula evaluated from: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html # # If the number of layers are changed this will need to be rewritten. - self.conv_end_w = self.conv2d_multi_layer(self.image_width, 3, kernel_size=3, padding=1, stride=2) self.conv_end_h = self.conv2d_multi_layer(self.image_height, 3, kernel_size=3, padding=1, stride=2) + self.conv_end_w = self.conv2d_multi_layer(self.image_width, 3, kernel_size=3, padding=1, stride=2) self._init_encoder() self._init_decoder() @@ -105,7 +105,7 @@ def _eval_decoder(self, x): x = self.dec_linear(x) x = x.reshape(x.shape[0], -1, self.conv_end_h, self.conv_end_w) x = self.decoder(x) - x = CenterCrop(size=(self.image_width, self.image_height))(x) + x = CenterCrop(size=(self.image_height, self.image_width))(x) return x def forward(self, batch): diff --git a/src/hyrax/models/hyrax_autoencoderv2.py b/src/hyrax/models/hyrax_autoencoderv2.py index 8101be440..eeea9ddea 100644 --- a/src/hyrax/models/hyrax_autoencoderv2.py +++ b/src/hyrax/models/hyrax_autoencoderv2.py @@ -8,6 +8,7 @@ # extra long import here to address a circular import issue from hyrax.models.model_registry import hyrax_model +from hyrax.models.model_utils import read_input_shape logger = logging.getLogger(__name__) @@ -27,18 +28,17 @@ class HyraxAutoencoderV2(nn.Module): V2 improvements: - Configurable final layer activation - Uses criterion and optimizer from config variables + + The layers are sized from ``[model.HyraxAutoencoderV2].input_shape``, the + ``[channels, height, width]`` of one input image. """ - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config - shape = data_sample.shape - logger.debug(f"Found shape: {shape} in data sample, using this to initialize model.") - - # Unpack the shape of the image (batch_size, num_channels, width, height) - # we'll ignore the batch_size during initialization. - _, self.num_input_channels, self.image_width, self.image_height = shape + self.input_shape = read_input_shape(config, "HyraxAutoencoderV2") + self.num_input_channels, self.image_height, self.image_width = self.input_shape self.c_hid = self.config["model"]["HyraxAutoencoderV2"]["base_channel_size"] self.latent_dim = self.config["model"]["HyraxAutoencoderV2"]["latent_dim"] @@ -47,8 +47,8 @@ def __init__(self, config, data_sample=None): # Formula evaluated from: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html # # If the number of layers are changed this will need to be rewritten. - self.conv_end_w = self.conv2d_multi_layer(self.image_width, 3, kernel_size=3, padding=1, stride=2) self.conv_end_h = self.conv2d_multi_layer(self.image_height, 3, kernel_size=3, padding=1, stride=2) + self.conv_end_w = self.conv2d_multi_layer(self.image_width, 3, kernel_size=3, padding=1, stride=2) self._init_encoder() self._init_decoder() @@ -129,7 +129,7 @@ def _eval_decoder(self, x): x = self.dec_linear(x) x = x.reshape(x.shape[0], -1, self.conv_end_h, self.conv_end_w) x = self.decoder(x) - x = CenterCrop(size=(self.image_width, self.image_height))(x) + x = CenterCrop(size=(self.image_height, self.image_width))(x) return x def forward(self, batch): diff --git a/src/hyrax/models/hyrax_cnn.py b/src/hyrax/models/hyrax_cnn.py index 61da78071..3fb4dd185 100644 --- a/src/hyrax/models/hyrax_cnn.py +++ b/src/hyrax/models/hyrax_cnn.py @@ -9,6 +9,7 @@ import torch.nn.functional as F # noqa N812 from .model_registry import hyrax_model +from .model_utils import read_input_shape logger = logging.getLogger(__name__) @@ -17,20 +18,17 @@ class HyraxCNN(nn.Module): """ This CNN is designed to work with datasets that are prepared with Hyrax's HSC Data Set class. + + The layers are sized from ``[model.HyraxCNN].input_shape``, the ``[channels, height, width]`` + of one input image. """ - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config - if data_sample is None: - raise ValueError("A `data_sample` must be provided to HyraxCNN for dynamic sizing.") - - image_sample = data_sample[0] - - # Unpack the shape of the image (batch_size, num_channels, width, height) - # we'll ignore the batch_size during initialization. - _, self.num_input_channels, self.image_width, self.image_height = image_sample.shape + self.input_shape = read_input_shape(config, "HyraxCNN") + self.num_input_channels, self.image_height, self.image_width = self.input_shape hidden_channels_1 = 6 hidden_channels_2 = 16 diff --git a/src/hyrax/models/hyrax_loopback.py b/src/hyrax/models/hyrax_loopback.py index db3883767..bd50f999d 100644 --- a/src/hyrax/models/hyrax_loopback.py +++ b/src/hyrax/models/hyrax_loopback.py @@ -11,7 +11,7 @@ class HyraxLoopback(nn.Module): """Simple model for testing which returns its own input""" - def __init__(self, config, data_sample=None): + def __init__(self, config): from functools import partial super().__init__() diff --git a/src/hyrax/models/hyrax_ts2vec.py b/src/hyrax/models/hyrax_ts2vec.py index 87fb2162f..e393c0371 100644 --- a/src/hyrax/models/hyrax_ts2vec.py +++ b/src/hyrax/models/hyrax_ts2vec.py @@ -364,8 +364,8 @@ class EncodingSettings: """The ``[model.HyraxTs2Vec]`` keys that describe the input rather than the architecture. Resolved and validated in one place so that :func:`build_event_sequence` and - :class:`HyraxTs2Vec` agree on the channel count, and so a bad setting is caught during - model pre-flight rather than midway through a run. + :class:`HyraxTs2Vec` agree on the channel count, and so a bad setting is caught when + the model is constructed rather than midway through a run. ``time_field``, ``flux_field`` and ``flux_err_field`` are flattened column names, as ``LSDBStreamDataset`` exposes nested columns. ``band_field`` is ``None`` for a @@ -684,7 +684,7 @@ class HyraxTs2Vec(nn.Module): dataset's ``max_sequence_length``. """ - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() # Must be set before __init__ returns: the @hyrax_model decorator reads self.config # immediately afterwards to build the criterion, optimizer, and scheduler. @@ -692,25 +692,10 @@ def __init__(self, config, data_sample=None): model_config = config["model"]["HyraxTs2Vec"] - if data_sample is None: - raise RuntimeError( - "HyraxTs2Vec needs a data_sample to size its input layer. Hyrax normally " - "supplies one from your dataset; when constructing the model by hand, pass " - "a (batch, time, channels) array." - ) - if getattr(data_sample, "ndim", None) != 3: - raise RuntimeError( - "HyraxTs2Vec expects a 3-dimensional (batch, time, channels) data sample, but " - f"got shape {getattr(data_sample, 'shape', type(data_sample))}. Hyrax builds " - "it with build_event_sequence(); pass an array of that shape when " - "constructing the model by hand." - ) - - self.input_dims = int(data_sample.shape[-1]) - logger.debug( - f"Found shape: {tuple(data_sample.shape)} in data sample, " - f"using input_dims={self.input_dims} to initialize model." - ) + # The input width is fixed by the encoding settings: flux, flux error and time gap, + # plus one one-hot channel per configured band. build_event_sequence() reads the + # same settings, so the data and the model agree by construction. + self.input_dims = EncodingSettings(config).n_channels self.output_dims = model_config["output_dims"] self.hidden_dims = model_config["hidden_dims"] diff --git a/src/hyrax/models/image_dcae.py b/src/hyrax/models/image_dcae.py index d1f8f793b..5600f77be 100644 --- a/src/hyrax/models/image_dcae.py +++ b/src/hyrax/models/image_dcae.py @@ -10,6 +10,7 @@ # extra long import here to address a circular import issue from hyrax.models.model_registry import hyrax_model +from hyrax.models.model_utils import read_input_shape class ArcsinhActivation(nn.Module): @@ -24,24 +25,17 @@ class ImageDCAE(nn.Module): """ This is an autoencoder with skipconnections that should work with arbitarily sized images with arbitrary number of channels. + + The bottleneck is sized from ``[model.ImageDCAE].input_shape``, the + ``[channels, height, width]`` of one input image. """ - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() - - if data_sample is None: - raise ValueError("data_sample must be provided to ImageDCAE for dynamic sizing.") - # Store input shape for dynamic sizing - self.input_shape = data_sample.shape self.config = config - # Extract dimensions from input shape - if len(self.input_shape) == 4: # Batch dimension included - self.num_input_channels = self.input_shape[1] - self.image_height, self.image_width = self.input_shape[2], self.input_shape[3] - else: # No batch dimension - self.num_input_channels = self.input_shape[0] - self.image_height, self.image_width = self.input_shape[1], self.input_shape[2] + self.input_shape = read_input_shape(config, "ImageDCAE") + self.num_input_channels, self.image_height, self.image_width = self.input_shape # Get latent dimension from config (similar to HyraxAutoencoder) self.latent_dim = config["model"]["ImageDCAE"]["latent_dim"] diff --git a/src/hyrax/models/model_utils.py b/src/hyrax/models/model_utils.py index 3e4631887..426a9d86c 100644 --- a/src/hyrax/models/model_utils.py +++ b/src/hyrax/models/model_utils.py @@ -57,3 +57,44 @@ def load_model_weights(config: dict, model, verb: str): msg = f"Model weights file {weights_file_path} did not load properly. Are you sure you are " msg += f"{verb}ing with the correct model?" raise RuntimeError(msg) from err + + +def read_input_shape(config: dict, model_name: str) -> tuple[int, int, int]: + """Read and validate ``[model.].input_shape`` from the configuration. + + Image models size their layers from this key at construction, so it has to describe + one input image completely: ``[channels, height, width]``. + + Parameters + ---------- + config : dict + Full runtime configuration + model_name : str + Name of the model's config subtable, e.g. ``"HyraxCNN"`` + + Returns + ------- + tuple[int, int, int] + ``(channels, height, width)`` + + Raises + ------ + ValueError + If the key is missing or is not a list of three positive integers. + """ + key = f"config['model']['{model_name}']['input_shape']" + shape = config["model"][model_name].get("input_shape") + + # `false` is the TOML "not set" sentinel, and bool is a subclass of int, so booleans + # must be rejected explicitly or `[3, true, 32]` would pass as `[3, 1, 32]`. + def _is_positive_int(value) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value > 0 + + if not isinstance(shape, (list, tuple)) or len(shape) != 3 or not all(map(_is_positive_int, shape)): + raise ValueError( + f"{key} must be [channels, height, width], three positive integers describing one " + f"input image, got {shape!r}. It has to match what the dataset produces (see " + "[data_set].crop_to)." + ) + + return int(shape[0]), int(shape[1]), int(shape[2]) diff --git a/src/hyrax/models/simclr.py b/src/hyrax/models/simclr.py index 81200d218..b18f47377 100644 --- a/src/hyrax/models/simclr.py +++ b/src/hyrax/models/simclr.py @@ -66,10 +66,9 @@ def __call__(self, x): class SimCLR(nn.Module): """SimCLR model. Implementation based on Chen, 2020""" - def __init__(self, config, shape): + def __init__(self, config): super().__init__() self.config = config - self.shape = shape proj_dim = config["model"]["SimCLR"]["projection_dimension"] temperature = config["model"]["SimCLR"]["temperature"] diff --git a/src/hyrax/pytorch_ignite.py b/src/hyrax/pytorch_ignite.py index 10e4b8ad4..7d84978c9 100644 --- a/src/hyrax/pytorch_ignite.py +++ b/src/hyrax/pytorch_ignite.py @@ -167,75 +167,21 @@ def _build_data_provider(config: dict, request: dict): return DataProvider(config, request) -def setup_model(config: dict, dataset: DataProvider) -> torch.nn.Module: +def setup_model(config: dict) -> torch.nn.Module: """Create a model object based on the configuration. Parameters ---------- config : dict The runtime configuration - dataset : DataProvider - The dataset object that will provide data to the model for training or - inference. Here it is only used to provide a data sample to the model so - that it can resize itself at runtime if necessary. Returns ------- torch.nn.Module An instance of the model class specified in the configuration """ - from hyrax.trace import reset_trace - - # Fetch model class specified in config and create an instance of it - model_cls = fetch_model_class(config) - - # Grab a single data sample - data_sample = dataset.sample_data() - - # Collate the data sample - collated_sample = dataset.collate([data_sample]) - - # Prepare the data sample with the model's prepare_inputs function - prepared_sample = model_cls.prepare_inputs(collated_sample) - - # Provide the sample for runtime modifications to the model architecture - retval = model_cls(config=config, data_sample=prepared_sample) # type: ignore[attr-defined] - - # After model pre-flighting succeeds (presumably) reset the trace so it represents - # just what the verb does afterward. - reset_trace() - - return retval - - -def setup_model_from_sample(config: dict, sample_batch: dict) -> torch.nn.Module: - """Create a model from a pre-formatted batch dict instead of a DataProvider. - - Like :func:`setup_model` but accepts a batch dict directly, bypassing the - DataLoader/DataProvider pipeline. Used by - :class:`~hyrax.verbs.infer_stream.InferStream` to pre-flight model - architecture without a dataset. - - Parameters - ---------- - config : dict - The runtime configuration. - sample_batch : dict - A representative batch with the same structure as batches that will be - processed later (e.g. ``{"object_id": [...], "data": {...}}``). - - Returns - ------- - torch.nn.Module - An instance of the model class specified in the configuration. - """ - from hyrax.trace import reset_trace - model_cls = fetch_model_class(config) - prepared_sample = model_cls.prepare_inputs(sample_batch) - model_instance = model_cls(config=config, data_sample=prepared_sample) - reset_trace() - return model_instance + return model_cls(config=config) # type: ignore[attr-defined] def dist_data_loader( diff --git a/src/hyrax/trace.py b/src/hyrax/trace.py index 56a3cdafe..b55b2eed8 100644 --- a/src/hyrax/trace.py +++ b/src/hyrax/trace.py @@ -120,9 +120,8 @@ def __enter__(self): self.config["train"]["epochs"] = 1 self.config["data_loader"]["batch_size"] = trace_batch_size - # Having cache running means some calls to getters aren't captured in a typical capture, - # because those same data were fetched during pre-flight, so we turn caching off on a - # trace to prevent this effect + # With the cache on, a getter asked for the same datum twice is only captured once, so + # caching is turned off on a trace to keep every call visible # # TODO: Rolling tracing. We will need a different solution here if we are ever running # the trace as a ring buffer, holding essentially the "last batch" up to a crash. @@ -146,17 +145,6 @@ def get_trace(): return trace_result -def reset_trace(): - """ - Reset the current global trace results object, removing all captured data - Valid to call whether trace is active. - """ - global trace_result - if trace_result: - logger.debug("Resetting Trace Results") - trace_result.reset() - - @dataclass class TraceDef: """ diff --git a/src/hyrax/verbs/infer.py b/src/hyrax/verbs/infer.py index 6af4eed26..9d175ed16 100644 --- a/src/hyrax/verbs/infer.py +++ b/src/hyrax/verbs/infer.py @@ -76,7 +76,7 @@ def run(self): if hasattr(provider, "data_cache"): provider.data_cache.disable() create_splits(config, dataset, persist=False) - model = setup_model(config, dataset["infer"]) + model = setup_model(config) logger.info( f"{Style.BRIGHT}{Fore.BLACK}{Back.GREEN}Inference model:{Style.RESET_ALL} " f"{model.__class__.__name__}" diff --git a/src/hyrax/verbs/infer_stream.py b/src/hyrax/verbs/infer_stream.py index 003ff4966..c877be220 100644 --- a/src/hyrax/verbs/infer_stream.py +++ b/src/hyrax/verbs/infer_stream.py @@ -33,43 +33,30 @@ def run_cli(self, args=None): "infer_stream is a programmatic API; use hyrax.infer_stream() in Python/notebook." ) - def run(self, sample_batch: dict | None = None) -> "InferStreamSession": + def run(self) -> "InferStreamSession": """Set up the model and return a session for streaming inference. There are two ways to drive the session: - 1. **Data-source driven** (``sample_batch=None``) — configure a streaming - dataset under ``[data_request.infer_stream]`` (e.g. ``KafkaStreamDataset``). - The model is pre-flighted from the stream itself and a DataLoader is built, - so the returned session can be iterated directly:: + 1. **Data-source driven** — configure a streaming dataset under + ``[data_request.infer_stream]`` (e.g. ``KafkaStreamDataset``). A DataLoader is + built over it, so the returned session can be iterated directly:: with hy.infer_stream() as session: for batch, results in session: ... - 2. **Manual** — pass a representative ``sample_batch`` and feed batches yourself:: + 2. **Manual** — leave ``[data_request.infer_stream]`` unset and feed batches + yourself:: - with hy.infer_stream(sample_batch=batch) as session: + with hy.infer_stream() as session: results = session.process(batch) - Parameters - ---------- - sample_batch : dict | None - A representative batch dict with ``"object_id"`` and model-specific data - fields, used to pre-flight the model architecture. When ``None``, the model - is pre-flighted from a ``[data_request.infer_stream]`` streaming dataset - instead. - Returns ------- InferStreamSession A context manager / session object. Iterate it (data-source driven) or call ``session.process(batch)`` (manual); call ``session.close()`` when done. - - Raises - ------ - ValueError - If ``sample_batch`` is None and no ``[data_request.infer_stream]`` is configured. """ from ignite.distributed import auto_model from ignite.distributed import device as idist_device @@ -83,7 +70,6 @@ def run(self, sample_batch: dict | None = None) -> "InferStreamSession": dist_data_loader, setup_dataset, setup_model, - setup_model_from_sample, ) from hyrax.tensorboardx_logger import close_tensorboard_logger, init_tensorboard_logger @@ -94,28 +80,16 @@ def run(self, sample_batch: dict | None = None) -> "InferStreamSession": # matching the ordering used by the train, infer and test verbs. results_dir = create_results_dir(config, "infer_stream") - # Build the model either from a configured streaming dataset (preferred, enables - # session iteration) or from an explicitly supplied sample batch. + model = setup_model(config) + + # A configured [data_request.infer_stream] group makes the session iterable over + # that stream. Without one the session is manual: the caller feeds process(). provider = None data_loader = None - if sample_batch is None: - if not config.get("data_request"): - raise ValueError( - "infer_stream requires either a `sample_batch` argument or a " - "[data_request.infer_stream] configuration to build the data source." - ) - datasets = setup_dataset(config, splits=InferStream.REQUIRED_DATA_GROUPS) - provider = datasets.get("infer_stream") - if provider is None: - raise ValueError( - "No [data_request.infer_stream] group found. Configure it with a " - "streaming dataset_class, or pass an explicit `sample_batch`." - ) - # Pre-flight the model from the stream (peeks one sample without losing it). - model = setup_model(config, provider) + if config.get("data_request"): + provider = setup_dataset(config, splits=InferStream.REQUIRED_DATA_GROUPS).get("infer_stream") + if provider is not None: data_loader = dist_data_loader(provider, config) - else: - model = setup_model_from_sample(config, sample_batch) # set model in eval mode model.eval() diff --git a/src/hyrax/verbs/test.py b/src/hyrax/verbs/test.py index 3b3921140..60acb6d0e 100644 --- a/src/hyrax/verbs/test.py +++ b/src/hyrax/verbs/test.py @@ -87,7 +87,7 @@ def run(self): if not isinstance(dataset, dict) or "test" not in dataset: raise RuntimeError("No test dataset available. Please configure a test dataset or split.") - model = setup_model(config, dataset["test"]) + model = setup_model(config) # Load model weights load_model_weights(config, model, "test") diff --git a/src/hyrax/verbs/to_onnx.py b/src/hyrax/verbs/to_onnx.py index 55d26b5a2..ca509cf40 100644 --- a/src/hyrax/verbs/to_onnx.py +++ b/src/hyrax/verbs/to_onnx.py @@ -96,10 +96,10 @@ def run(self, input_model_directory: str = None): if not weights_file_path.exists(): raise RuntimeError(f"Could not find trained model weights: {weights_file_path}") - # Use the config in the model directory to load the dataset(s) and create - # The data loader instance to provide a data sample to the ONNX exporter. + # Use the config in the model directory to build the model and load the dataset(s); + # the dataset provides the sample batch the ONNX exporter traces with. dataset = setup_dataset(config_from_training, splits=("infer",), shuffle=False) - model = setup_model(config_from_training, dataset["infer"]) + model = setup_model(config_from_training) # Load the trained weights and send the model to the CPU for ONNX export. model.load(weights_file_path) model.train(False) diff --git a/src/hyrax/verbs/train.py b/src/hyrax/verbs/train.py index 7819d81e8..c805dc280 100644 --- a/src/hyrax/verbs/train.py +++ b/src/hyrax/verbs/train.py @@ -86,7 +86,7 @@ def run(self): splits=Train.REQUIRED_DATA_GROUPS + Train.OPTIONAL_DATA_GROUPS, ) create_splits(config, dataset, results_dir=results_dir, persist=True) - model = setup_model(config, dataset["train"]) + model = setup_model(config) # we separate out the cases here for two reasons: # 1. GitHub merge checks don't pass with idist.Parallel diff --git a/src/hyrax/verbs/train_stream.py b/src/hyrax/verbs/train_stream.py index 876c366ca..f28f27eac 100644 --- a/src/hyrax/verbs/train_stream.py +++ b/src/hyrax/verbs/train_stream.py @@ -54,7 +54,7 @@ def run(self) -> "TrainStreamSession": Raises ------ ValueError - If ``sample_batch`` is None and no ``[data_request.train_stream]`` is configured. + If no ``[data_request.train_stream]`` is configured. """ from pathlib import Path @@ -86,11 +86,11 @@ def run(self) -> "TrainStreamSession": # Create a timestamped results directory and start logging. results_dir = create_results_dir(config, "train_stream") - # From this point on, the provider/data stream is live (setup_model peeks a sample). - # If any subsequent setup step raises, no caller will get a session to call close() - # on, so we must tear down the provider ourselves to avoid leaking the connection. + # The provider may already hold resources (e.g. an LSDB dataset's dask client). If + # any subsequent setup step raises, no caller will get a session to call close() + # on, so we must tear down the provider ourselves to avoid leaking them. try: - model = setup_model(config, provider) + model = setup_model(config) data_loader = dist_data_loader(provider, config) # Put the model in training mode, and place it on the correct device diff --git a/tests/hyrax/conftest.py b/tests/hyrax/conftest.py index d12f087c2..258284c35 100644 --- a/tests/hyrax/conftest.py +++ b/tests/hyrax/conftest.py @@ -279,7 +279,7 @@ class ContextWritingLoopback(nn.Module): and the model registry perform. """ - def __init__(self, config, data_sample=None): + def __init__(self, config): from functools import partial super().__init__() diff --git a/tests/hyrax/test_config_migrations/test_migration_007.py b/tests/hyrax/test_config_migrations/test_migration_007.py new file mode 100644 index 000000000..fc3a8a4e8 --- /dev/null +++ b/tests/hyrax/test_config_migrations/test_migration_007.py @@ -0,0 +1,34 @@ +"""Tests for migration 007: removal of KafkaStreamDataset.peek_timeout.""" + +import tomlkit + +from hyrax.config_migrations import CURRENT_CONFIG_VERSION, migrate_config + + +def test_migrate_007_removes_peek_timeout(): + """peek_timeout is dropped and the rest of the Kafka table is untouched.""" + cfg = tomlkit.parse( + "config_version = 7\n" + "[data_set.KafkaStreamDataset]\n" + 'topics = ["alerts"]\n' + "peek_timeout = 120.0\n" + "poll_timeout = 1.0\n" + ) + + migrated = migrate_config(cfg) + + kafka = migrated["data_set"]["KafkaStreamDataset"] + assert "peek_timeout" not in kafka + assert kafka["topics"] == ["alerts"] + assert kafka["poll_timeout"] == 1.0 + assert migrated["config_version"] == CURRENT_CONFIG_VERSION + + +def test_migrate_007_without_kafka_table_is_noop(): + """A v7 config that never configured Kafka passes through cleanly.""" + cfg = tomlkit.parse("config_version = 7\n[data_set]\nseed = 42\n") + + migrated = migrate_config(cfg) + + assert migrated["data_set"]["seed"] == 42 + assert migrated["config_version"] == CURRENT_CONFIG_VERSION diff --git a/tests/hyrax/test_csv_dataset.py b/tests/hyrax/test_csv_dataset.py index 8de9c2e06..070ccc7fd 100644 --- a/tests/hyrax/test_csv_dataset.py +++ b/tests/hyrax/test_csv_dataset.py @@ -53,26 +53,3 @@ def test_csv_dataset_column_getters(test_hyrax_csv_dataset): assert csv_dataset.get_object_id(0) == 1001 assert csv_dataset.get_ra(0) == 30.5 assert csv_dataset.get_classification(0) == "star" - - -def test_csv_dataset_sample_data(test_hyrax_csv_dataset): - """Check that sample_data returns the first row correctly""" - dataset = test_hyrax_csv_dataset.prepare() - - # Get the underlying HyraxCSVDataset instance - csv_dataset = dataset["train"]._primary_or_first_dataset() - sample = csv_dataset.sample_data() - - # Check that sample has the expected structure - assert "data" in sample - assert "object_id" in sample["data"] - assert "ra" in sample["data"] - assert "dec" in sample["data"] - assert "magnitude" in sample["data"] - assert "flux" in sample["data"] - assert "classification" in sample["data"] - - # Check that sample values match the first row - assert sample["data"]["object_id"] == 1001 - assert sample["data"]["ra"] == 30.5 - assert sample["data"]["classification"] == "star" diff --git a/tests/hyrax/test_data_provider.py b/tests/hyrax/test_data_provider.py index 522fdf499..e057a3cfb 100644 --- a/tests/hyrax/test_data_provider.py +++ b/tests/hyrax/test_data_provider.py @@ -381,86 +381,6 @@ def test_metadata_fields_with_friendly_name(data_provider): assert field in all_fields -def test_sample_data(): - """Test that sample_data returns a dictionary with the expected - structure. - - We don't use the test fixture here so that this can be a little more - flexible and self-contained. - - - The expected result structure is: - { - 'random_0': { - 'object_id': , - 'image': array(...), - 'label': 'cat' - }, - 'random_1': { - 'image': array(...) - }, - 'object_id': - } - """ - from hyrax import Hyrax - - h = Hyrax() - - multimodal_config = { - "random_0": { - "dataset_class": "HyraxRandomDataset", - "data_directory": "./in_memory_0", - "fields": ["object_id", "image", "label"], - "dataset_config": { - "HyraxRandomDataset": { - "shape": [2, 16, 16], - }, - }, - "primary_id_field": "object_id", - }, - "random_1": { - "dataset_class": "HyraxRandomDataset", - "data_directory": "./in_memory_1", - "fields": ["image"], - "dataset_config": { - "HyraxRandomDataset": { - "shape": [5, 16, 16], - "seed": 4200, - }, - }, - }, - } - - h.config["data_request"] = multimodal_config - dp = DataProvider(h.config, multimodal_config) - dp.prepare_datasets() - - sample = dp.sample_data() - - assert isinstance(sample, dict) - assert "random_0" in sample - assert "random_1" in sample - assert "object_id" in sample - assert len(sample) == 3 - assert isinstance(sample["random_0"], dict) - assert isinstance(sample["random_1"], dict) - assert len(sample["random_0"]) == 3 - assert len(sample["random_1"]) == 1 - - for friendly_name in ["random_0", "random_1"]: - dataset_sample = sample[friendly_name] - assert "image" in dataset_sample - - if friendly_name == "random_0": - assert "object_id" in dataset_sample - assert "label" in dataset_sample - - # Verify the dataset_config overrides actually took effect - # (default shape is [2, 5, 5] — these should differ) - assert sample["random_0"]["image"].shape == (2, 16, 16) - assert sample["random_1"]["image"].shape == (5, 16, 16) - - def test_data_provider_get_item(data_provider): """Basic test to ensure that different index values return different data.""" dp = data_provider diff --git a/tests/hyrax/test_e2e.py b/tests/hyrax/test_e2e.py index 07f474a01..bfc7f1b74 100644 --- a/tests/hyrax/test_e2e.py +++ b/tests/hyrax/test_e2e.py @@ -72,10 +72,14 @@ def hyrax_instance(tmp_dataset_path, dataset_spec, model_class_name, tmp_path): h.config["general"]["data_dir"] = str(tmp_dataset_path) h.config["general"]["results_dir"] = str(tmp_path) - # Configure dataset-specific settings + # Configure dataset-specific settings. The hsc1k demo data has three filters and + # cutouts of at least 96x96 pixels, so crop every image to that and size the model + # to match. HyraxCifarDataset matches the model's default input_shape as shipped. + if sample_data == "hsc1k": + h.config["data_set"]["crop_to"] = [96, 96] + h.config["model"][model_class_name]["input_shape"] = [3, 96, 96] if dataset_class_name == "FitsImageDataset" and sample_data == "hsc1k": h.config["data_set"]["filter_catalog"] = str(tmp_dataset_path / "manifest.fits") - h.config["data_set"]["crop_to"] = [100, 100] h.config["model"]["name"] = model_class_name h.config["train"]["epochs"] = 1 diff --git a/tests/hyrax/test_hyrax_autoencoderv2.py b/tests/hyrax/test_hyrax_autoencoderv2.py index 4f73f526a..4585efd56 100644 --- a/tests/hyrax/test_hyrax_autoencoderv2.py +++ b/tests/hyrax/test_hyrax_autoencoderv2.py @@ -10,7 +10,7 @@ # trims the reconstruction back down to the original size. 20x20 comfortably survives # the three encoder downsamples to a nonzero spatial size, matching the image size used # in the HyraxCNN tests. -IMAGE_SHAPE = (1, 20, 20) # (channels, width, height) +IMAGE_SHAPE = (1, 20, 20) # (channels, height, width) BASE_CHANNEL_SIZE = 4 LATENT_DIM = 8 @@ -22,6 +22,7 @@ def make_model_config(base_channel_size=BASE_CHANNEL_SIZE, latent_dim=LATENT_DIM return { "model": { "HyraxAutoencoderV2": { + "input_shape": list(IMAGE_SHAPE), "base_channel_size": base_channel_size, "latent_dim": latent_dim, "final_layer": "tanh", @@ -34,12 +35,6 @@ def make_model_config(base_channel_size=BASE_CHANNEL_SIZE, latent_dim=LATENT_DIM } -def make_data_sample(): - """A stand-in for the `data_sample` HyraxAutoencoderV2 uses at init time to size its - layers dynamically. Only the image tensor's shape is inspected.""" - return torch.zeros(BATCH_SIZE, *IMAGE_SHAPE) - - def make_hyrax_batch(batch_size=BATCH_SIZE, seed=0): """Build a batch in Hyrax's collated "dict of lists" format: a single dict of keys where each value is a numpy array holding every datapoint @@ -69,7 +64,7 @@ def to_tensor_batch(image_array): @pytest.fixture def model(): """A freshly constructed HyraxAutoencoderV2 sized for 1x20x20 images.""" - return HyraxAutoencoderV2(config=make_model_config(), data_sample=make_data_sample()) + return HyraxAutoencoderV2(config=make_model_config()) def test_prepare_inputs_extracts_image(): @@ -162,11 +157,11 @@ def test_train_batch_is_reproducible_with_fixed_seed(): batch = to_tensor_batch(HyraxAutoencoderV2.prepare_inputs(hyrax_batch)) torch.manual_seed(1234) - model_a = HyraxAutoencoderV2(config=make_model_config(), data_sample=make_data_sample()) + model_a = HyraxAutoencoderV2(config=make_model_config()) loss_a = model_a.train_batch(batch)["loss"] torch.manual_seed(1234) - model_b = HyraxAutoencoderV2(config=make_model_config(), data_sample=make_data_sample()) + model_b = HyraxAutoencoderV2(config=make_model_config()) loss_b = model_b.train_batch(batch)["loss"] np.testing.assert_almost_equal(loss_a, loss_b) @@ -187,6 +182,7 @@ def hyrax_autoencoderv2_pipeline(tmp_path_factory): h.config["model"]["name"] = "HyraxAutoencoderV2" h.config["model"]["HyraxAutoencoderV2"]["base_channel_size"] = BASE_CHANNEL_SIZE h.config["model"]["HyraxAutoencoderV2"]["latent_dim"] = LATENT_DIM + h.config["model"]["HyraxAutoencoderV2"]["input_shape"] = list(IMAGE_SHAPE) h.config["criterion"]["name"] = "torch.nn.MSELoss" h.config["train"]["epochs"] = 1 h.config["data_loader"]["batch_size"] = BATCH_SIZE diff --git a/tests/hyrax/test_hyrax_cnn.py b/tests/hyrax/test_hyrax_cnn.py index 67097d170..d595a34e1 100644 --- a/tests/hyrax/test_hyrax_cnn.py +++ b/tests/hyrax/test_hyrax_cnn.py @@ -9,7 +9,7 @@ # stages, which requires at least an 18x18 input to avoid collapsing to a 0x0 # tensor before the final layers. 20x20 is the smallest round size above that # floor, so we use it here instead of a truly minimal test image. -IMAGE_SHAPE = (1, 20, 20) # (channels, width, height) +IMAGE_SHAPE = (1, 20, 20) # (channels, height, width) OUTPUT_CLASSES = 2 @@ -18,7 +18,7 @@ def make_model_config(output_classes=OUTPUT_CLASSES): `hyrax_model` to build its optimizer/criterion/scheduler directly, without going through a full `hyrax.Hyrax()` instance.""" return { - "model": {"HyraxCNN": {"output_classes": output_classes}}, + "model": {"HyraxCNN": {"input_shape": list(IMAGE_SHAPE), "output_classes": output_classes}}, "criterion": {"name": "torch.nn.CrossEntropyLoss"}, "optimizer": {"name": "torch.optim.SGD"}, "torch.optim.SGD": {"lr": 0.01, "momentum": 0.9}, @@ -26,12 +26,6 @@ def make_model_config(output_classes=OUTPUT_CLASSES): } -def make_data_sample(): - """A stand-in for the `data_sample` HyraxCNN uses at init time to size its - layers dynamically. Only the image tensor's shape is inspected.""" - return (torch.zeros(BATCH_SIZE, *IMAGE_SHAPE), torch.zeros(BATCH_SIZE, dtype=torch.long)) - - def make_hyrax_batch(batch_size=BATCH_SIZE, seed=0): """Build a batch in Hyrax's collated "dict of lists" format: a single dict of keys where each value is a numpy array holding every datapoint @@ -65,7 +59,7 @@ def to_tensor_batch(prepared_inputs): def model(): """A freshly constructed HyraxCNN sized for 1x20x20 images and a binary (0/1) label, matching the dataset format described for this model.""" - return HyraxCNN(config=make_model_config(), data_sample=make_data_sample()) + return HyraxCNN(config=make_model_config()) def test_prepare_inputs_extracts_image_and_label(): @@ -155,11 +149,11 @@ def test_train_batch_is_reproducible_with_fixed_seed(): batch = to_tensor_batch(HyraxCNN.prepare_inputs(hyrax_batch)) torch.manual_seed(1234) - model_a = HyraxCNN(config=make_model_config(), data_sample=make_data_sample()) + model_a = HyraxCNN(config=make_model_config()) loss_a = model_a.train_batch(batch)["loss"] torch.manual_seed(1234) - model_b = HyraxCNN(config=make_model_config(), data_sample=make_data_sample()) + model_b = HyraxCNN(config=make_model_config()) loss_b = model_b.train_batch(batch)["loss"] np.testing.assert_almost_equal(loss_a, loss_b) @@ -179,6 +173,7 @@ def hyrax_cnn_pipeline(tmp_path_factory): h = hyrax.Hyrax() h.config["model"]["name"] = "HyraxCNN" h.config["model"]["HyraxCNN"]["output_classes"] = OUTPUT_CLASSES + h.config["model"]["HyraxCNN"]["input_shape"] = list(IMAGE_SHAPE) h.config["train"]["epochs"] = 1 h.config["data_loader"]["batch_size"] = BATCH_SIZE h.config["general"]["results_dir"] = str(results_dir) diff --git a/tests/hyrax/test_hyrax_ts2vec.py b/tests/hyrax/test_hyrax_ts2vec.py index 68f9554f6..9ca46ce9e 100644 --- a/tests/hyrax/test_hyrax_ts2vec.py +++ b/tests/hyrax/test_hyrax_ts2vec.py @@ -450,13 +450,14 @@ def test_prepare_inputs_outside_a_verb_run_explains_itself(): HyraxTs2Vec.prepare_inputs(batch) -def test_model_sizes_its_input_layer_from_the_prepared_sample(hyrax_config): - """setup_model runs prepare_inputs before construction, so data_sample is the series.""" +def test_model_sizes_its_input_layer_from_the_configured_bands(hyrax_config): + """The input width comes from the encoding settings, and build_event_sequence agrees with it.""" samples = [light_curve(np.arange(8.0), bands=["g"] * 8)] - data_sample = build_event_sequence(collate(samples), hyrax_config) - model = HyraxTs2Vec(hyrax_config, data_sample=data_sample) + series = build_event_sequence(collate(samples), hyrax_config) + model = HyraxTs2Vec(hyrax_config) assert model.input_dims == N_LEADING_CHANNELS + len(hyrax_config["model"]["HyraxTs2Vec"]["bands"]) + assert model.input_dims == series.shape[-1] def test_train_batch_runs_on_an_encoded_batch(hyrax_config): @@ -473,7 +474,7 @@ def test_train_batch_runs_on_an_encoded_batch(hyrax_config): for _ in range(4) ] series = build_event_sequence(collate(samples), hyrax_config) - model = HyraxTs2Vec(hyrax_config, data_sample=series) + model = HyraxTs2Vec(hyrax_config) metrics = model.train_batch(torch.from_numpy(series)) @@ -486,7 +487,7 @@ def test_infer_batch_keeps_every_row(hyrax_config): samples = [light_curve(np.arange(8.0), bands=["g"] * 8), light_curve([], bands=[])] series = build_event_sequence(collate(samples), hyrax_config) - model = HyraxTs2Vec(hyrax_config, data_sample=series) + model = HyraxTs2Vec(hyrax_config) model.eval() with torch.no_grad(): diff --git a/tests/hyrax/test_infer_stream.py b/tests/hyrax/test_infer_stream.py index 5775c07c7..100e99727 100644 --- a/tests/hyrax/test_infer_stream.py +++ b/tests/hyrax/test_infer_stream.py @@ -124,11 +124,18 @@ def make(self): return created -def test_infer_stream_requires_sample_or_request(): - """With no sample_batch and no data_request, run() raises a clear error.""" - h = hyrax.Hyrax() - with pytest.raises(ValueError, match="sample_batch"): - h.infer_stream() +def test_infer_stream_without_data_request_is_manual(tmp_path): + """With no [data_request.infer_stream], run() returns a manual session: fed by process(), not iterable.""" + h = _build_manual_config(tmp_path, "HyraxLoopback") + batch = _image_batch(["id0", "id1"], start=0) + + with h.infer_stream() as session: + assert session.data_loader is None + with pytest.raises(RuntimeError, match="no data source"): + list(session) + results = session.process(batch) + + np.testing.assert_array_equal(results, batch["data"]["image"]) def test_infer_stream_iterates_streaming_dataset(tmp_path, monkeypatch): @@ -174,7 +181,7 @@ def test_infer_stream_iterates_streaming_dataset(tmp_path, monkeypatch): # Loopback returns its image input; one result row per object. assert results.shape[0] == len(ids) - # The peeked sample (used for model pre-flighting) is not lost. + # Every streamed sample reaches the model; none is consumed during setup. assert sorted(seen_ids) == ["id0", "id1", "id2"] @@ -183,7 +190,7 @@ def test_process_returns_array_for_tensor_model(tmp_path): h = _build_manual_config(tmp_path, "HyraxLoopback") batch = _image_batch(["id0", "id1"], start=0) - with h.infer_stream(sample_batch=batch) as session: + with h.infer_stream() as session: results = session.process(batch) assert isinstance(results, np.ndarray) @@ -195,7 +202,7 @@ def test_process_returns_dict_for_dict_model(tmp_path): h = _build_manual_config(tmp_path, "DictReturnLoopback") batches = [_image_batch(["id0", "id1"], start=0), _image_batch(["id2"], start=6)] - with h.infer_stream(sample_batch=batches[0]) as session: + with h.infer_stream() as session: for batch in batches: results = session.process(batch) assert isinstance(results, dict) diff --git a/tests/hyrax/test_input_shape.py b/tests/hyrax/test_input_shape.py new file mode 100644 index 000000000..6fa4ee773 --- /dev/null +++ b/tests/hyrax/test_input_shape.py @@ -0,0 +1,77 @@ +"""Tests for how the built-in image models size themselves from ``[model.].input_shape``.""" + +import copy + +import pytest +import torch + +import hyrax +from hyrax.models.hyrax_autoencoder import HyraxAutoencoder +from hyrax.models.hyrax_autoencoderv2 import HyraxAutoencoderV2 +from hyrax.models.hyrax_cnn import HyraxCNN +from hyrax.models.image_dcae import ImageDCAE +from hyrax.models.model_utils import read_input_shape + +IMAGE_MODELS = [ + pytest.param("HyraxCNN", HyraxCNN, id="HyraxCNN"), + pytest.param("HyraxAutoencoder", HyraxAutoencoder, id="HyraxAutoencoder"), + pytest.param("HyraxAutoencoderV2", HyraxAutoencoderV2, id="HyraxAutoencoderV2"), + pytest.param("ImageDCAE", ImageDCAE, id="ImageDCAE"), +] + +BAD_SHAPES = [ + pytest.param(False, id="false-sentinel"), + pytest.param([3, 32], id="too-short"), + pytest.param([1, 3, 32, 32], id="too-long"), + pytest.param([3, "32", 32], id="string"), + pytest.param([3, 0, 32], id="zero"), + pytest.param([3, -32, 32], id="negative"), + pytest.param([3, True, 32], id="bool"), + pytest.param(32, id="scalar"), +] + + +@pytest.fixture(scope="module") +def default_config(): + """The default Hyrax configuration, loaded once for the module.""" + return hyrax.Hyrax().config + + +@pytest.fixture +def config(default_config): + """A fresh copy of the default configuration for each test to mutate.""" + return copy.deepcopy(default_config) + + +def _batch_for(model_name, shape, batch_size=2): + """A batch of zeros in the form the model's forward expects.""" + images = torch.zeros(batch_size, *shape) + if model_name == "HyraxCNN": + return images, torch.zeros(batch_size, dtype=torch.long) + return images + + +@pytest.mark.parametrize("model_name,model_cls", IMAGE_MODELS) +def test_default_input_shape_builds_the_model(config, model_name, model_cls): + """Every image model constructs from the shipped default config and accepts matching images.""" + model = model_cls(config) + + assert model.input_shape == (3, 32, 32) + assert model(_batch_for(model_name, (3, 32, 32))).shape[0] == 2 + + +@pytest.mark.parametrize("model_name,model_cls", IMAGE_MODELS) +@pytest.mark.parametrize("bad_shape", BAD_SHAPES) +def test_bad_input_shape_raises_naming_the_key(config, model_name, model_cls, bad_shape): + """A malformed input_shape fails at construction with the config key in the message.""" + config["model"][model_name]["input_shape"] = bad_shape + + with pytest.raises(ValueError, match=rf"config\['model'\]\['{model_name}'\]\['input_shape'\]"): + model_cls(config) + + +def test_read_input_shape_accepts_a_tuple(config): + """A tuple, as tests and notebooks tend to write, is as good as a TOML list.""" + config["model"]["HyraxCNN"]["input_shape"] = (1, 20, 20) + + assert read_input_shape(config, "HyraxCNN") == (1, 20, 20) diff --git a/tests/hyrax/test_kafka_stream_dataset.py b/tests/hyrax/test_kafka_stream_dataset.py index 93c8ea7a3..a4168a2b6 100644 --- a/tests/hyrax/test_kafka_stream_dataset.py +++ b/tests/hyrax/test_kafka_stream_dataset.py @@ -272,67 +272,6 @@ def test_close_does_not_commit(monkeypatch): assert consumer.commits == [] -def test_peek_sample_closes_consumer_on_interrupt(monkeypatch): - """An interrupted peek must not leave a subscribed consumer behind. - - This is the notebook case: a KeyboardInterrupt out of peek_sample used to leave a live - consumer holding the topic's partitions, so the next run joined the same group and was - assigned nothing. - """ - dataset = _build_dataset() - consumer = _patch_consumer(monkeypatch, dataset, [], poll_error=KeyboardInterrupt()) - - with pytest.raises(KeyboardInterrupt): - dataset.peek_sample() - - assert consumer.closed - assert dataset._consumer is None - - -def test_peek_sample_closes_consumer_on_error(monkeypatch): - """Any exception out of peek_sample releases the consumer.""" - dataset = _build_dataset() - consumer = _patch_consumer(monkeypatch, dataset, [], poll_error=RuntimeError("broker exploded")) - - with pytest.raises(RuntimeError, match="broker exploded"): - dataset.peek_sample() - - assert consumer.closed - assert dataset._consumer is None - - -def test_peek_timeout_raises_with_diagnostic(monkeypatch): - """peek_sample gives up instead of blocking forever, and says what it was waiting on.""" - h = hyrax.Hyrax() - ds_config = h.config["data_set"]["KafkaStreamDataset"] - ds_config["topics"] = "test-topic" - ds_config["peek_timeout"] = 0.05 - dataset = KafkaStreamDataset(h.config) - consumer = _patch_consumer(monkeypatch, dataset, []) - - with pytest.raises(TimeoutError) as excinfo: - dataset.peek_sample() - - message = str(excinfo.value) - assert "peek_timeout" in message - assert "test-topic" in message - # An empty assignment is the signature of a consumer left over from an earlier run, - # so the error has to point at the command that reveals it. - assert "kafka-consumer-groups" in message - assert consumer.closed - - -def test_peek_timeout_disabled_by_false(monkeypatch): - """The TOML `false` sentinel restores the unbounded wait.""" - h = hyrax.Hyrax() - ds_config = h.config["data_set"]["KafkaStreamDataset"] - ds_config["topics"] = "test-topic" - ds_config["peek_timeout"] = False - dataset = KafkaStreamDataset(h.config) - - assert dataset.peek_timeout is None - - def test_offsets_committed_only_after_batch_is_consumed(monkeypatch): """Offsets advance when the caller comes back for more, not when a batch is handed out.""" dataset = _build_dataset(batch_size=2, batch_flush_timeout=0.0) @@ -375,22 +314,6 @@ def test_auto_commit_disabled(): assert dataset.consumer_config["enable.auto.offset.store"] is True -def test_peek_sample_buffers_and_replays(monkeypatch): - """peek_sample returns a flat sample that is not lost; it leads the first batch.""" - dataset = _build_dataset(batch_size=5, batch_flush_timeout=0.0) - messages = [_make_message("first", [[1.0]]), _make_message("second", [[2.0]])] - _patch_consumer(monkeypatch, dataset, messages, stop_when_exhausted=True) - - peeked = dataset.peek_sample() - assert peeked["object_id"] == "first" - - batches = list(dataset) - - assert len(batches) == 1 - # The peeked message is replayed as the first sample of the first batch. - assert [s["object_id"] for s in batches[0]] == ["first", "second"] - - def test_extra_credentials(tmp_path): """Extra credentials are passed to the consumer constructor.""" h = hyrax.Hyrax() diff --git a/tests/hyrax/test_list_registries.py b/tests/hyrax/test_list_registries.py index 0cbe5a4fa..a00313fb2 100644 --- a/tests/hyrax/test_list_registries.py +++ b/tests/hyrax/test_list_registries.py @@ -19,7 +19,7 @@ def test_list_models_includes_registered_model(): @hyrax_model class ZZZTestListModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config diff --git a/tests/hyrax/test_lsdb_stream_dataset.py b/tests/hyrax/test_lsdb_stream_dataset.py index 6cf625aac..5dc63da96 100644 --- a/tests/hyrax/test_lsdb_stream_dataset.py +++ b/tests/hyrax/test_lsdb_stream_dataset.py @@ -307,38 +307,6 @@ def test_empty_chunk_is_tolerated(monkeypatch, tiny_catalog): assert [len(batch) for batch in batches] == [4, 2] -def test_peek_sample_does_not_lose_a_row(monkeypatch, tiny_catalog): - """The peeked row is replayed at the head of the first batch.""" - dataset = _build_dataset(tiny_catalog, batch_size=4) - _with_chunks(monkeypatch, dataset, [_chunk(0, 3), _chunk(3, 3)]) - - sample = dataset.peek_sample() - assert sample["object_id"] == "id00" - - batches = list(dataset) - assert _ids(batches) == [f"id{i:02d}" for i in range(6)] - - -def test_peek_sample_twice_returns_distinct_rows_and_both_replay(monkeypatch, tiny_catalog): - """Repeated peeks advance a cursor rather than re-reading the same row.""" - dataset = _build_dataset(tiny_catalog, batch_size=4) - _with_chunks(monkeypatch, dataset, [_chunk(0, 3), _chunk(3, 3)]) - - assert dataset.peek_sample()["object_id"] == "id00" - assert dataset.peek_sample()["object_id"] == "id01" - - assert _ids(list(dataset)) == [f"id{i:02d}" for i in range(6)] - - -def test_peek_sample_on_empty_stream_raises(monkeypatch, tiny_catalog): - """A catalog with no rows fails with an actionable error, not StopIteration.""" - dataset = _build_dataset(tiny_catalog) - _with_chunks(monkeypatch, dataset, []) - - with pytest.raises(RuntimeError, match="produced no rows"): - dataset.peek_sample() - - def test_stop_ends_iteration_and_flushes_partial(monkeypatch, tiny_catalog): """stop() lets the current chunk finish, then flushes what is buffered.""" dataset = _build_dataset(tiny_catalog, batch_size=4) diff --git a/tests/hyrax/test_model_registry.py b/tests/hyrax/test_model_registry.py index 619dfe7aa..fb848ee96 100644 --- a/tests/hyrax/test_model_registry.py +++ b/tests/hyrax/test_model_registry.py @@ -11,7 +11,7 @@ def test_use_model_optimizer(): @hyrax_model class TestModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -32,7 +32,7 @@ def test_use_config_optimizer(): @hyrax_model class TestModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -51,7 +51,7 @@ def test_no_optimizer_defined_logs_warning(caplog): @hyrax_model class TestModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -71,7 +71,7 @@ def test_use_model_criterion(): @hyrax_model class TestModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -91,7 +91,7 @@ def test_use_config_criterion(): @hyrax_model class TestModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -110,7 +110,7 @@ def test_no_criterion_defined_logs_warning(caplog): @hyrax_model class TestModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -129,7 +129,7 @@ def test_criterion_defined_in_model_and_config(caplog): @hyrax_model class TestModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -151,7 +151,7 @@ def test_optimizer_defined_in_model_and_config(caplog): @hyrax_model class TestModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -174,7 +174,7 @@ def test_use_model_scheduler(): @hyrax_model class TestModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -194,7 +194,7 @@ def test_use_config_scheduler(caplog): @hyrax_model class TestModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -213,7 +213,7 @@ def test_no_scheduler_defined_logs_warning(caplog): @hyrax_model class TestModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -232,7 +232,7 @@ def test_scheduler_defined_in_model_and_config(caplog): @hyrax_model class TestModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) diff --git a/tests/hyrax/test_patch_model.py b/tests/hyrax/test_patch_model.py index 67d42573c..4e8c2a8a1 100644 --- a/tests/hyrax/test_patch_model.py +++ b/tests/hyrax/test_patch_model.py @@ -8,7 +8,7 @@ class DummyModelOne(nn.Module): """A dummy model used to test patching of static methods like prepare_inputs""" - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() # The optimizer needs at least one weight, so we add a dummy module here self.unused_module = nn.Linear(1, 1) @@ -25,7 +25,7 @@ class DummyModelTwo(nn.Module): """A dummy model used to test patching, that uses the default prepare_inputs method by default.""" - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() # The optimizer needs at least one weight, so we add a dummy module here self.unused_module = nn.Linear(1, 1) @@ -51,7 +51,7 @@ def test_patch_prepare_inputs(tmp_path): } # create an instance of the dummy model - model = DummyModelOne(config=config, data_sample=None) + model = DummyModelOne(config=config) # manually update the prepare_inputs static method to be something simple # don't wrap this with staticmethod(...) because that would be a double wrapping. @@ -64,7 +64,7 @@ def test_patch_prepare_inputs(tmp_path): assert (tmp_path / "prepare_inputs.py").exists() # create a new instance of the dummy model and call .load() with the correct path - new_model = DummyModelOne(config=config, data_sample=None) + new_model = DummyModelOne(config=config) # verify that the new model's prepare_inputs method is the default one input_data = 3.0 @@ -92,7 +92,7 @@ def test_patch_prepare_inputs_over_default(tmp_path): } # create an instance of the dummy model - model = DummyModelTwo(config=config, data_sample=None) + model = DummyModelTwo(config=config) # manually update the prepare_inputs static method to be something simple # don't wrap this with staticmethod(...) because that would be a double wrapping. @@ -105,7 +105,7 @@ def test_patch_prepare_inputs_over_default(tmp_path): assert (tmp_path / "prepare_inputs.py").exists() # create a new instance of the dummy model and call .load() with the correct path - new_model = DummyModelTwo(config=config, data_sample=None) + new_model = DummyModelTwo(config=config) # verify that the new model's prepare_inputs method is the default one input_data = {"data": {"image": 3}} diff --git a/tests/hyrax/test_plugin_utils.py b/tests/hyrax/test_plugin_utils.py index 136a82983..589f5cf47 100644 --- a/tests/hyrax/test_plugin_utils.py +++ b/tests/hyrax/test_plugin_utils.py @@ -157,7 +157,7 @@ def test_torch_load_with_map_location(tmp_path): # Create a simple model @hyrax_model class SimpleModel(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.linear = nn.Linear(10, 5) @@ -218,7 +218,7 @@ def test_torch_load_keeps_optimizer_bound_to_parameters(tmp_path): @hyrax_model class TinyNet(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.linear = nn.Linear(4, 1) @@ -271,7 +271,7 @@ def test_torch_load_only_calls_to_tensor_if_prepare_inputs_missing(tmp_path, cap # Create a simple model with prepare_inputs @hyrax_model class SimpleModelWithPrepareInputs(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.linear = nn.Linear(10, 5) @@ -337,7 +337,7 @@ def test_torch_load_calls_to_tensor_when_prepare_inputs_missing(tmp_path, caplog # Create a simple model @hyrax_model class SimpleModelNoInputs(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.linear = nn.Linear(10, 5) diff --git a/tests/hyrax/test_prepare_inputs_migration.py b/tests/hyrax/test_prepare_inputs_migration.py index 2b4c2c587..a6ad08dec 100644 --- a/tests/hyrax/test_prepare_inputs_migration.py +++ b/tests/hyrax/test_prepare_inputs_migration.py @@ -14,7 +14,7 @@ def test_model_with_only_to_tensor(): @hyrax_model class TestModelWithToTensor(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -52,7 +52,7 @@ def test_model_without_prepare_inputs_or_to_tensor(): @hyrax_model class TestModelNoInputMethod(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -87,7 +87,7 @@ def test_model_without_prepare_inputs_or_to_tensor_raises_error(): @hyrax_model class TestModelNoInputMethod(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -127,7 +127,7 @@ def test_save_and_load_model_with_to_tensor_warns(tmp_path, caplog): @hyrax_model class TestModelWithToTensor(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -191,7 +191,7 @@ def test_load_model_with_old_to_tensor_file(tmp_path, caplog): @hyrax_model class TestModelForLoading(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) @@ -247,7 +247,7 @@ def test_load_model_with_no_serialized_input_functions(tmp_path, caplog): @hyrax_model class TestModelForLoading(nn.Module): - def __init__(self, config, data_sample=None): + def __init__(self, config): super().__init__() self.config = config self.unused_module = nn.Linear(1, 1) diff --git a/tests/hyrax/test_streaming_data_provider.py b/tests/hyrax/test_streaming_data_provider.py index 59bb3f879..f68f57870 100644 --- a/tests/hyrax/test_streaming_data_provider.py +++ b/tests/hyrax/test_streaming_data_provider.py @@ -1,4 +1,4 @@ -"""Tests for StreamingDataProvider — structuring, collation, sample_data, and routing. +"""Tests for StreamingDataProvider — structuring, collation, and routing. Reuses the FakeConsumer pattern from ``test_kafka_stream_dataset.py``: the provider builds a KafkaStreamDataset internally, so tests patch ``provider._stream._make_consumer``. @@ -158,21 +158,6 @@ def test_collate_matches_infer_contract(): assert collated["data"]["image"].shape == (4, 1, 2) -def test_sample_data_returns_structured_and_is_not_lost(monkeypatch): - """sample_data peeks one structured sample that still appears in the first batch.""" - provider = _build_provider(batch_size=5, flush=0.0, fields=("image",)) - messages = [_make_message("first", [[1.0]]), _make_message("second", [[2.0]])] - _patch_stream(monkeypatch, provider, messages) - - sample = provider.sample_data() - assert sample["object_id"] == "first" - assert sample["data"]["image"].shape == (1, 1) - - batches = list(provider) - assert len(batches) == 1 - assert [s["object_id"] for s in batches[0]] == ["first", "second"] - - def test_dist_data_loader_end_to_end(monkeypatch): """The provider flows through dist_data_loader and yields collated batch dicts.""" provider = _build_provider(batch_size=2, fields=("image", "flux")) diff --git a/tests/hyrax/test_train_stream.py b/tests/hyrax/test_train_stream.py index 52b8dd7e7..b27ebfb0a 100644 --- a/tests/hyrax/test_train_stream.py +++ b/tests/hyrax/test_train_stream.py @@ -149,7 +149,7 @@ def test_train_stream_iterates_streaming_dataset(tmp_path, monkeypatch): assert "loss" in metrics results_dir = session._results_dir - # The peeked sample (used for model pre-flighting) is not lost. + # Every streamed sample reaches the model; none is consumed during setup. assert sorted(seen_ids) == ["id0", "id1", "id2"] # Final weights were persisted on close. assert (results_dir / "example_model.pth").exists() From 90305d4315ef318c9d23a24bc9d6bd5a1b2a9fe8 Mon Sep 17 00:00:00 2001 From: Drew Date: Wed, 28 Oct 2026 12:26:33 -0700 Subject: [PATCH 2/3] Fixing a spot that was missed in trace example notebook. --- docs/pre_executed/trace_feature.ipynb | 667 +++++++++++++++++++++++++- 1 file changed, 643 insertions(+), 24 deletions(-) diff --git a/docs/pre_executed/trace_feature.ipynb b/docs/pre_executed/trace_feature.ipynb index a41b80e57..6de6a49d6 100644 --- a/docs/pre_executed/trace_feature.ipynb +++ b/docs/pre_executed/trace_feature.ipynb @@ -32,7 +32,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "metadata": {}, "outputs": [], "source": [ @@ -54,7 +54,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "metadata": {}, "outputs": [], "source": [ @@ -75,11 +75,12 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "h.config[\"model\"][\"name\"] = \"HyraxAutoencoder\"\n", + "h.config[\"model\"][\"HyraxAutoencoder\"][\"input_shape\"] = [1, 32, 32]\n", "\n", "# Use a small 1-channel 32×32 image shape to keep things fast\n", "h.config[\"data_set\"][\"HyraxRandomDataset\"][\"shape\"] = [1, 32, 32]\n", @@ -109,7 +110,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 4, "metadata": {}, "outputs": [], "source": [ @@ -163,9 +164,66 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 5, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "[2026-10-28 12:25:02,960 hyrax.trace:WARNING] Starting Trace\n", + "[2026-10-28 12:25:02,961 hyrax.trace:WARNING] Trace mode enabled, will only run a single batch of length 2\n", + "[2026-10-28 12:25:02,979 hyrax.models.model_registry:INFO] Setting model's self.optimizer from config: torch.optim.SGD with arguments: {'lr': 0.01, 'momentum': 0.9}.\n", + "[2026-10-28 12:25:02,979 hyrax.models.model_registry:INFO] Setting model's self.criterion from config: torch.nn.CrossEntropyLoss with default arguments.\n", + "[2026-10-28 12:25:02,980 hyrax.models.model_registry:INFO] Setting model's self.scheduler from config: torch.optim.lr_scheduler.ExponentialLR\n", + "with arguments: {'gamma': 1}.\n", + "[2026-10-28 12:25:03,045 hyrax.verbs.train:INFO] Using a single process for training.\n", + "[2026-10-28 12:25:03,045 hyrax.verbs.train:INFO] \u001b[1m\u001b[30m\u001b[42mTraining model:\u001b[0m HyraxAutoencoder\n", + "[2026-10-28 12:25:03,046 hyrax.verbs.train:INFO] \u001b[1m\u001b[30m\u001b[42mTraining dataset(s):\u001b[0m\n", + "{'train': Name: data (primary dataset)\n", + " Dataset class: HyraxRandomDataset\n", + " Data location: /home/drew/code/hyrax/docs/pre_executed/trace_data\n", + " Selected items: 50\n", + " Primary ID field: object_id\n", + " Requested fields: image\n", + "}\n", + "2026-10-28 12:25:03,091 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset ', 'batch_size': 2, 'collate_fn': , 'pin_memory': True}\n", + "2026/10/28 12:25:03 INFO mlflow.system_metrics.system_metrics_monitor: Skip logging GPU metrics. Set logger level to DEBUG for more details.\n", + "2026/10/28 12:25:03 INFO mlflow.system_metrics.system_metrics_monitor: Started monitoring system metrics.\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "e1f81166e8824e2eabf3ce0b54da732f", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "100%|##########| 1/1 [00:00 image duration=0.00381 ms\n", + "\t\tinputs:\n", + "\t\t\tindex = 6\n", + "\t\toutputs:\n", + "\t\t\timage = \n", + "\t\t\n", + "\t\tdata__get_object_id(index) -> object_id duration=0.00378 ms\n", + "\t\tinputs:\n", + "\t\t\tindex = 6\n", + "\t\toutputs:\n", + "\t\t\tobject_id = '29'\n", + "\t\t\n", + "\t\tdata__get_image(index) -> image duration=0.00177 ms\n", + "\t\tinputs:\n", + "\t\t\tindex = 16\n", + "\t\toutputs:\n", + "\t\t\timage = \n", + "\t\t\n", + "\t\tdata__get_object_id(index) -> object_id duration=0.000736 ms\n", + "\t\tinputs:\n", + "\t\t\tindex = 16\n", + "\t\toutputs:\n", + "\t\t\tobject_id = '39'\n", + "\t\t\n", + "\t]\n", + "\tresolve_data: [\n", + "\t\tDataProvider__resolve_data(index) -> data_dict duration=0.034 ms\n", + "\t\tinputs:\n", + "\t\t\tindex = 6\n", + "\t\toutputs:\n", + "\t\t\tdata_dict = {\n", + "\t\t\t\tdata: {\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t} \n", + "\t\t\t\tobject_id: '29' \n", + "\t\t\t}\n", + "\t\t\n", + "\t\tDataProvider__resolve_data(index) -> data_dict duration=0.01 ms\n", + "\t\tinputs:\n", + "\t\t\tindex = 16\n", + "\t\toutputs:\n", + "\t\t\tdata_dict = {\n", + "\t\t\t\tdata: {\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t} \n", + "\t\t\t\tobject_id: '39' \n", + "\t\t\t}\n", + "\t\t\n", + "\t]\n", + "\tfield_level_collation: []\n", + "\tcollate: [\n", + "\t\tDataProvider__collate(batch_dicts) -> batch_dict duration=0.124 ms\n", + "\t\tinputs:\n", + "\t\t\tbatch_dicts = [\n", + "\t\t\t\t{\n", + "\t\t\t\t\tdata: {\n", + "\t\t\t\t\t\timage: \n", + "\t\t\t\t\t} \n", + "\t\t\t\t\tobject_id: '29' \n", + "\t\t\t\t}\n", + "\t\t\t\t{\n", + "\t\t\t\t\tdata: {\n", + "\t\t\t\t\t\timage: \n", + "\t\t\t\t\t} \n", + "\t\t\t\t\tobject_id: '39' \n", + "\t\t\t\t}\n", + "\t\t\t]\n", + "\t\toutputs:\n", + "\t\t\tbatch_dict = {\n", + "\t\t\t\tobject_id: \n", + "\t\t\t\tdata: {\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t\tmask: \n", + "\t\t\t\t} \n", + "\t\t\t}\n", + "\t\t\n", + "\t\tdata__collate(samples) -> batch_dict duration=0.0376 ms\n", + "\t\tinputs:\n", + "\t\t\tsamples = [\n", + "\t\t\t\t{\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t}\n", + "\t\t\t\t{\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t}\n", + "\t\t\t]\n", + "\t\toutputs:\n", + "\t\t\tbatch_dict = {\n", + "\t\t\t\timage: \n", + "\t\t\t\tmask: \n", + "\t\t\t}\n", + "\t\t\n", + "\t\tDataProvider__handle_nans(batch_dict) -> batch_dict_no_nan duration=0.0665 ms\n", + "\t\tinputs:\n", + "\t\t\tbatch_dict = {\n", + "\t\t\t\tobject_id: \n", + "\t\t\t\tdata: {\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t\tmask: \n", + "\t\t\t\t} \n", + "\t\t\t}\n", + "\t\toutputs:\n", + "\t\t\tbatch_dict_no_nan = {\n", + "\t\t\t\tobject_id: \n", + "\t\t\t\tdata: {\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t\tmask: \n", + "\t\t\t\t} \n", + "\t\t\t}\n", + "\t\t\n", + "\t]\n", + "\tprepare_inputs: [\n", + "\t\tHyraxAutoencoder__prepare_inputs(batch_dict) -> batch_ndarray duration=0.00623 ms\n", + "\t\tinputs:\n", + "\t\t\tbatch_dict = {\n", + "\t\t\t\tobject_id: \n", + "\t\t\t\tdata: {\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t\tmask: \n", + "\t\t\t\t} \n", + "\t\t\t}\n", + "\t\toutputs:\n", + "\t\t\tbatch_ndarray = \n", + "\t\t\n", + "\t]\n", + "\tevaluation: [\n", + "\t\tHyraxAutoencoder__train_batch(batch) -> loss_dict duration=968 ms\n", + "\t\tinputs:\n", + "\t\t\tbatch = \n", + "\t\toutputs:\n", + "\t\t\tloss_dict = {\n", + "\t\t\t\tloss: \n", + "\t\t\t}\n", + "\t\t\n", + "\t\tHyraxAutoencoder__forward(batch) -> batch_results duration=504 ms\n", + "\t\tinputs:\n", + "\t\t\tbatch = \n", + "\t\toutputs:\n", + "\t\t\tbatch_results = \n", + "\t\t\n", + "\t]\n", + "}\n", + "\n" + ] + } + ], "source": [ "print(trace_result)" ] @@ -221,9 +431,77 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 7, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[\n", + "\tDataProvider__collate(batch_dicts) -> batch_dict duration=0.124 ms\n", + "\tinputs:\n", + "\t\tbatch_dicts = [\n", + "\t\t\t{\n", + "\t\t\t\tdata: {\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t} \n", + "\t\t\t\tobject_id: '29' \n", + "\t\t\t}\n", + "\t\t\t{\n", + "\t\t\t\tdata: {\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t} \n", + "\t\t\t\tobject_id: '39' \n", + "\t\t\t}\n", + "\t\t]\n", + "\toutputs:\n", + "\t\tbatch_dict = {\n", + "\t\t\tobject_id: \n", + "\t\t\tdata: {\n", + "\t\t\t\timage: \n", + "\t\t\t\tmask: \n", + "\t\t\t} \n", + "\t\t}\n", + "\t\n", + "\tdata__collate(samples) -> batch_dict duration=0.0376 ms\n", + "\tinputs:\n", + "\t\tsamples = [\n", + "\t\t\t{\n", + "\t\t\t\timage: \n", + "\t\t\t}\n", + "\t\t\t{\n", + "\t\t\t\timage: \n", + "\t\t\t}\n", + "\t\t]\n", + "\toutputs:\n", + "\t\tbatch_dict = {\n", + "\t\t\timage: \n", + "\t\t\tmask: \n", + "\t\t}\n", + "\t\n", + "\tDataProvider__handle_nans(batch_dict) -> batch_dict_no_nan duration=0.0665 ms\n", + "\tinputs:\n", + "\t\tbatch_dict = {\n", + "\t\t\tobject_id: \n", + "\t\t\tdata: {\n", + "\t\t\t\timage: \n", + "\t\t\t\tmask: \n", + "\t\t\t} \n", + "\t\t}\n", + "\toutputs:\n", + "\t\tbatch_dict_no_nan = {\n", + "\t\t\tobject_id: \n", + "\t\t\tdata: {\n", + "\t\t\t\timage: \n", + "\t\t\t\tmask: \n", + "\t\t\t} \n", + "\t\t}\n", + "\t\n", + "]\n" + ] + } + ], "source": [ "# Attribute-style access\n", "collate_stage = trace_result.collate\n", @@ -232,9 +510,32 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 8, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[\n", + "\tHyraxAutoencoder__train_batch(batch) -> loss_dict duration=968 ms\n", + "\tinputs:\n", + "\t\tbatch = \n", + "\toutputs:\n", + "\t\tloss_dict = {\n", + "\t\t\tloss: \n", + "\t\t}\n", + "\t\n", + "\tHyraxAutoencoder__forward(batch) -> batch_results duration=504 ms\n", + "\tinputs:\n", + "\t\tbatch = \n", + "\toutputs:\n", + "\t\tbatch_results = \n", + "\t\n", + "]\n" + ] + } + ], "source": [ "# Dictionary-style access — identical result\n", "evaluation_stage = trace_result[\"evaluation\"]\n", @@ -251,9 +552,19 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 9, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "resolve_data calls : 2\n", + "collate calls : 3\n", + "evaluation calls : 2\n" + ] + } + ], "source": [ "print(f\"resolve_data calls : {len(trace_result.resolve_data)}\")\n", "print(f\"collate calls : {len(trace_result.collate)}\")\n", @@ -272,9 +583,41 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 10, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "DataProvider__collate(batch_dicts) -> batch_dict duration=0.124 ms\n", + "inputs:\n", + "\tbatch_dicts = [\n", + "\t\t{\n", + "\t\t\tdata: {\n", + "\t\t\t\timage: \n", + "\t\t\t} \n", + "\t\t\tobject_id: '29' \n", + "\t\t}\n", + "\t\t{\n", + "\t\t\tdata: {\n", + "\t\t\t\timage: \n", + "\t\t\t} \n", + "\t\t\tobject_id: '39' \n", + "\t\t}\n", + "\t]\n", + "outputs:\n", + "\tbatch_dict = {\n", + "\t\tobject_id: \n", + "\t\tdata: {\n", + "\t\t\timage: \n", + "\t\t\tmask: \n", + "\t\t} \n", + "\t}\n", + "\n" + ] + } + ], "source": [ "# Get the first call in the collate stage\n", "first_collate_call = trace_result.collate[0]\n", @@ -283,9 +626,61 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 11, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "{'object_id': array(['29', '39'], dtype=' data_dict duration=0.034 ms\n", + "inputs:\n", + "\tindex = 6\n", + "outputs:\n", + "\tdata_dict = {\n", + "\t\tdata: {\n", + "\t\t\timage: \n", + "\t\t} \n", + "\t\tobject_id: '29' \n", + "\t}\n", + "\n" + ] + } + ], "source": [ "# Get a list of all calls to a particular function by providing the display name\n", "# (visible from the print output above)\n", @@ -317,9 +731,214 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 13, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "[2026-10-28 12:25:05,049 hyrax.trace:WARNING] Starting Trace\n", + "[2026-10-28 12:25:05,050 hyrax.trace:WARNING] Trace mode enabled, will only run a single batch of length 2\n", + "[2026-10-28 12:25:05,067 hyrax.models.model_registry:INFO] Setting model's self.optimizer from config: torch.optim.SGD with arguments: {'lr': 0.01, 'momentum': 0.9}.\n", + "[2026-10-28 12:25:05,068 hyrax.models.model_registry:INFO] Setting model's self.criterion from config: torch.nn.CrossEntropyLoss with default arguments.\n", + "[2026-10-28 12:25:05,068 hyrax.models.model_registry:INFO] Setting model's self.scheduler from config: torch.optim.lr_scheduler.ExponentialLR\n", + "with arguments: {'gamma': 1}.\n", + "[2026-10-28 12:25:05,068 hyrax.verbs.infer:INFO] \u001b[1m\u001b[30m\u001b[42mInference model:\u001b[0m HyraxAutoencoder\n", + "[2026-10-28 12:25:05,069 hyrax.verbs.infer:INFO] \u001b[1m\u001b[30m\u001b[42mInference dataset(s):\u001b[0m\n", + "{'infer': Name: data (primary dataset)\n", + " Dataset class: HyraxRandomDataset\n", + " Data location: /home/drew/code/hyrax/docs/pre_executed/trace_data\n", + " Selected items: 50\n", + " Primary ID field: object_id\n", + " Requested fields: image\n", + "}\n", + "2026-10-28 12:25:05,070 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset ', 'pin_memory': True}\n", + "[2026-10-28 12:25:05,080 hyrax.models.model_utils:INFO] Updated config['infer']['model_weights_file'] to: /home/drew/code/hyrax/docs/pre_executed/results/20261028-122502-train-tWy2/example_model.pth\n", + "[2026-10-28 12:25:05,081 hyrax.verbs.infer:INFO] Saving inference results at: /home/drew/code/hyrax/docs/pre_executed/results/20261028-122505-infer-umzD\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "e3ebefd6ccd74b518148af752bcea562", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "100%|##########| 1/1 [00:00 image duration=0.00595 ms\n", + "\t\tinputs:\n", + "\t\t\tindex = 0\n", + "\t\toutputs:\n", + "\t\t\timage = \n", + "\t\t\n", + "\t\tdata__get_object_id(index) -> object_id duration=0.00508 ms\n", + "\t\tinputs:\n", + "\t\t\tindex = 0\n", + "\t\toutputs:\n", + "\t\t\tobject_id = '23'\n", + "\t\t\n", + "\t\tdata__get_image(index) -> image duration=0.000414 ms\n", + "\t\tinputs:\n", + "\t\t\tindex = 1\n", + "\t\toutputs:\n", + "\t\t\timage = \n", + "\t\t\n", + "\t\tdata__get_object_id(index) -> object_id duration=0.000404 ms\n", + "\t\tinputs:\n", + "\t\t\tindex = 1\n", + "\t\toutputs:\n", + "\t\t\tobject_id = '24'\n", + "\t\t\n", + "\t]\n", + "\tresolve_data: [\n", + "\t\tDataProvider__resolve_data(index) -> data_dict duration=0.0248 ms\n", + "\t\tinputs:\n", + "\t\t\tindex = 0\n", + "\t\toutputs:\n", + "\t\t\tdata_dict = {\n", + "\t\t\t\tdata: {\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t} \n", + "\t\t\t\tobject_id: '23' \n", + "\t\t\t}\n", + "\t\t\n", + "\t\tDataProvider__resolve_data(index) -> data_dict duration=0.00515 ms\n", + "\t\tinputs:\n", + "\t\t\tindex = 1\n", + "\t\toutputs:\n", + "\t\t\tdata_dict = {\n", + "\t\t\t\tdata: {\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t} \n", + "\t\t\t\tobject_id: '24' \n", + "\t\t\t}\n", + "\t\t\n", + "\t]\n", + "\tfield_level_collation: []\n", + "\tcollate: [\n", + "\t\tDataProvider__collate(batch_dicts) -> batch_dict duration=0.11 ms\n", + "\t\tinputs:\n", + "\t\t\tbatch_dicts = [\n", + "\t\t\t\t{\n", + "\t\t\t\t\tdata: {\n", + "\t\t\t\t\t\timage: \n", + "\t\t\t\t\t} \n", + "\t\t\t\t\tobject_id: '23' \n", + "\t\t\t\t}\n", + "\t\t\t\t{\n", + "\t\t\t\t\tdata: {\n", + "\t\t\t\t\t\timage: \n", + "\t\t\t\t\t} \n", + "\t\t\t\t\tobject_id: '24' \n", + "\t\t\t\t}\n", + "\t\t\t]\n", + "\t\toutputs:\n", + "\t\t\tbatch_dict = {\n", + "\t\t\t\tobject_id: \n", + "\t\t\t\tdata: {\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t\tmask: \n", + "\t\t\t\t} \n", + "\t\t\t}\n", + "\t\t\n", + "\t\tdata__collate(samples) -> batch_dict duration=0.0362 ms\n", + "\t\tinputs:\n", + "\t\t\tsamples = [\n", + "\t\t\t\t{\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t}\n", + "\t\t\t\t{\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t}\n", + "\t\t\t]\n", + "\t\toutputs:\n", + "\t\t\tbatch_dict = {\n", + "\t\t\t\timage: \n", + "\t\t\t\tmask: \n", + "\t\t\t}\n", + "\t\t\n", + "\t\tDataProvider__handle_nans(batch_dict) -> batch_dict_no_nan duration=0.0523 ms\n", + "\t\tinputs:\n", + "\t\t\tbatch_dict = {\n", + "\t\t\t\tobject_id: \n", + "\t\t\t\tdata: {\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t\tmask: \n", + "\t\t\t\t} \n", + "\t\t\t}\n", + "\t\toutputs:\n", + "\t\t\tbatch_dict_no_nan = {\n", + "\t\t\t\tobject_id: \n", + "\t\t\t\tdata: {\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t\tmask: \n", + "\t\t\t\t} \n", + "\t\t\t}\n", + "\t\t\n", + "\t]\n", + "\tprepare_inputs: [\n", + "\t\tHyraxAutoencoder_inst_prepare_inputs(batch_dict) -> batch_ndarray duration=0.00256 ms\n", + "\t\tinputs:\n", + "\t\t\tbatch_dict = {\n", + "\t\t\t\tobject_id: \n", + "\t\t\t\tdata: {\n", + "\t\t\t\t\timage: \n", + "\t\t\t\t\tmask: \n", + "\t\t\t\t} \n", + "\t\t\t}\n", + "\t\toutputs:\n", + "\t\t\tbatch_ndarray = \n", + "\t\t\n", + "\t]\n", + "\tevaluation: [\n", + "\t\tHyraxAutoencoder__infer_batch(batch) -> batch_results duration=0.525 ms\n", + "\t\tinputs:\n", + "\t\t\tbatch = \n", + "\t\toutputs:\n", + "\t\t\tbatch_results = \n", + "\t\t\n", + "\t\tHyraxAutoencoder__forward(batch) -> batch_results duration=0.519 ms\n", + "\t\tinputs:\n", + "\t\t\tbatch = \n", + "\t\toutputs:\n", + "\t\t\tbatch_results = \n", + "\t\t\n", + "\t]\n", + "}\n", + "\n" + ] + } + ], "source": [ "infer_trace = h.infer(trace=2)\n", "print(infer_trace)" @@ -358,7 +977,7 @@ ], "metadata": { "kernelspec": { - "display_name": "kbmod", + "display_name": "hyrax", "language": "python", "name": "python3" }, @@ -372,7 +991,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.11" + "version": "3.13.15" } }, "nbformat": 4, From 4a80b47073ae9791db92bcaae945e250b3b15155 Mon Sep 17 00:00:00 2001 From: Drew Oldag <47493171+drewoldag@users.noreply.github.com> Date: Wed, 28 Oct 2026 14:14:52 -0700 Subject: [PATCH 3/3] Update src/hyrax/models/model_utils.py --- src/hyrax/models/model_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hyrax/models/model_utils.py b/src/hyrax/models/model_utils.py index 426a9d86c..59f0bc7dd 100644 --- a/src/hyrax/models/model_utils.py +++ b/src/hyrax/models/model_utils.py @@ -83,7 +83,7 @@ def read_input_shape(config: dict, model_name: str) -> tuple[int, int, int]: If the key is missing or is not a list of three positive integers. """ key = f"config['model']['{model_name}']['input_shape']" - shape = config["model"][model_name].get("input_shape") + shape = config["model"][model_name]["input_shape"] # `false` is the TOML "not set" sentinel, and bool is a subclass of int, so booleans # must be rejected explicitly or `[3, true, 32]` would pass as `[3, 1, 32]`.