diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst index 47fe681..19b63b7 100644 --- a/docs/source/api/index.rst +++ b/docs/source/api/index.rst @@ -126,6 +126,7 @@ Signature rtol: float = 1e-3, atol: float = 1e-5, strict: bool = True, + graph_configs: GraphConfigs | None = None, ) Parameters @@ -147,6 +148,8 @@ Parameters ``torch.allclose()``. Defaults to ``1e-5``. * ``strict`` (bool, optional): Raise on validation failure. Defaults to ``True``. +* ``graph_configs`` (GraphConfigs, optional): Graph-level metadata to serialize + under the generated YAML ``pipeline.configs`` entry. Behavior ~~~~~~~~ @@ -161,6 +164,42 @@ The method performs the complete pipeline: #. Generate ``{name}.png`` if ``visualize=True``. #. Log graph statistics. +``GraphConfigs`` +---------------- + +Graph-level metadata passed to :func:`~leapp.compile_graph`. + +.. code-block:: python + + from leapp import GraphConfigs + + leapp.compile_graph( + graph_configs=GraphConfigs( + frequency_hz=50, + extra={"runtime": "isaac_lab"}, + ), + ) + +This emits fields under the generated ``pipeline.configs`` entry: + +.. code-block:: yaml + + pipeline: + data_flow: {} + feedback_flow: {} + inputs: {} + outputs: {} + configs: + frequency: 50 + runtime: isaac_lab + +Parameters +~~~~~~~~~~ + +* ``frequency`` (float | int, optional): Graph-level execution frequency in Hz. +* ``extra`` (dict, optional): Additional graph-level fields. Keys are flattened + into the ``pipeline.configs`` entry rather than serialized under an ``extra`` key. + Generated artifacts ~~~~~~~~~~~~~~~~~~~ @@ -204,7 +243,7 @@ Output YAML structure python version: "3.12.9" torch version: "2.7.0+cu126" leapp version: "0.5.2" - leapp config version: "1.1" + leapp config version: "1.2" cuda version: "12.6" os: Linux @@ -222,6 +261,34 @@ The method logs statistics including: Semantic metadata ================= +``TemporalAxis`` +-------------------- + +Mark one ``element_names`` axis as temporal and attach an absolute period in +milliseconds to the tensor YAML entry. + +.. code-block:: python + + from leapp import TensorSemantics, TemporalAxis + + TensorSemantics( + name="actions", + ref=actions, + element_names=[TemporalAxis(period_ms=100), ["hip", "knee", "ankle"]], + ) + +This emits the reserved bare axis sentinel and the period metadata: + +.. code-block:: yaml + + element_names: + - __temporal_axis__ + - [hip, knee, ankle] + temporal_period_ms: 100 + +``__temporal_axis__`` is reserved for LEAPP output; use +``TemporalAxis`` rather than writing the sentinel directly. + ``TensorSemantics`` ------------------- @@ -252,7 +319,8 @@ Parameters the tensor. Enum values are serialized to their string ``.value``. * ``element_names`` (list | str, optional): Human-readable element names. A string becomes ``[[name]]``; a flat list of strings becomes ``[[...]]``; a - per-dimension list is preserved with optional ``None`` entries. + per-dimension list is preserved with optional ``None`` entries. Include + ``TemporalAxis(...)`` as a bare axis item to mark that axis temporal. * ``extra`` (dict, optional): Additional semantic fields. Keys are flattened into the tensor YAML entry rather than serialized under an ``extra`` key. diff --git a/docs/source/guides/semantics.rst b/docs/source/guides/semantics.rst index 0c53d9c..beee0ea 100644 --- a/docs/source/guides/semantics.rst +++ b/docs/source/guides/semantics.rst @@ -295,6 +295,55 @@ automatically: TensorSemantics("gravity", tensor, element_names="z") +``TemporalAxis`` +-------------------- + +Use ``TemporalAxis`` inside ``element_names`` to mark one tensor axis as +temporal and record the period between samples on that axis. This is useful for +chunked outputs such as an action tensor shaped ``[num_chunks, num_joints]``. + +Place ``TemporalAxis`` directly in the outer ``element_names`` list. Do not +wrap it in a list. LEAPP serializes the temporal axis as the reserved +``__temporal_axis__`` sentinel and emits ``temporal_period_ms`` as a sibling +field on the tensor entry. + +.. code-block:: python + + from leapp import TemporalAxis + + TensorSemantics( + "actions", + actions, + kind=OutputKindEnum.JOINT_TORQUES, + element_names=[ + TemporalAxis(period_ms=100), + ["hip", "knee", "ankle"], + ], + ) + +.. code-block:: yaml + + - name: actions + dtype: float32 + shape: [4, 3] + type: tensor + kind: target/joint/torques + element_names: + - __temporal_axis__ + - [hip, knee, ankle] + temporal_period_ms: 100 + +Downstream consumers can find the temporal axis by locating +``__temporal_axis__`` in ``element_names``. LEAPP does not validate +``temporal_period_ms`` against ``GraphConfigs.frequency``. + +.. warning:: + + ``__temporal_axis__`` is reserved for LEAPP output. Use + ``TemporalAxis`` in Python annotations rather than writing the + sentinel directly. A tensor may contain at most one temporal axis marker. + + ``extra`` --------- @@ -401,10 +450,48 @@ Limitations within the same node's inputs (or outputs). Duplicate names raise an error. #. **Semantic fields are optional** --- all semantic fields (``kind``, - ``element_names``, and ``extra``) are optional. A ``TensorSemantics`` - with no semantic fields behaves identically to passing a raw tensor with - the same name. + ``element_names``, temporal metadata from ``TemporalAxis``, and + ``extra``) are optional. A ``TensorSemantics`` with no semantic fields + behaves identically to passing a raw tensor with the same name. #. **Extra fields are flattened** --- keys in ``extra`` become top-level YAML fields on the tensor entry. Avoid keys that collide with built-in fields - such as ``name``, ``dtype``, ``shape``, ``type``, ``kind``, or - ``element_names``. + such as ``name``, ``dtype``, ``shape``, ``type``, ``kind``, + ``element_names``, ``temporal_period_ms``, or ``__temporal_axis__``. + +Graph-level semantics +===================== + +Use ``GraphConfigs`` for metadata that applies to the whole exported LEAPP +bundle rather than to a specific tensor. ``GraphConfigs`` is passed to +:func:`~leapp.compile_graph`, after tracing has stopped. + +.. code-block:: python + + from leapp import GraphConfigs + + leapp.compile_graph( + graph_configs=GraphConfigs( + frequency=50, + extra={"skip-first-run": True}, + ), + ) + +``frequency`` describes the graph-level execution frequency in Hz. ``extra`` +works like ``TensorSemantics.extra`` within ``configs``: keys are flattened into +the generated ``pipeline.configs`` entry rather than nested under an ``extra`` +key. + +.. code-block:: yaml + + pipeline: + data_flow: {} + feedback_flow: {} + inputs: {} + outputs: {} + configs: + frequency: 50 + skip-first-run: true + +Graph-level semantics are independent of tensor-level temporal metadata. LEAPP +does not validate ``TemporalAxis`` values against ``GraphConfigs.frequency``. + diff --git a/docs/source/spelling_wordlist.txt b/docs/source/spelling_wordlist.txt index 3745ac7..dff647b 100644 --- a/docs/source/spelling_wordlist.txt +++ b/docs/source/spelling_wordlist.txt @@ -107,6 +107,7 @@ TracedTensor TracedTensors TracedTensorNode TensorSemantics +TemporalAxis InferenceManager InputKindEnum OutputKindEnum diff --git a/leapp/__init__.py b/leapp/__init__.py index fff22a4..9d3d9b5 100644 --- a/leapp/__init__.py +++ b/leapp/__init__.py @@ -28,10 +28,10 @@ from .inference_manager import InferenceManager from .leapp import annotate, start, stop, compile_graph from .utils.enums import InputKindEnum, OutputKindEnum -from .utils.tensor_description import TensorSemantics +from .utils.tensor_description import GraphConfigs, TensorSemantics, TemporalAxis __version__ = "0.5.2" -__config_version__ = "1.1" +__config_version__ = "1.2" __author__ = "Frank Lai" __email__ = "frlai@nvidia.com" @@ -46,4 +46,6 @@ "compile_graph", "__version__", "TensorSemantics", + "TemporalAxis", + "GraphConfigs", ] diff --git a/leapp/leapp.py b/leapp/leapp.py index 8f5f9ab..02f9b84 100644 --- a/leapp/leapp.py +++ b/leapp/leapp.py @@ -28,12 +28,14 @@ from .utils.tracing_lock import TracingLock from .utils.utils import get_system_info from .export_manager import ExportManager +from .utils.tensor_description import GraphConfigs _MANAGER = ExportManager() -def start(name, save_path=".", verbose=False, dry_run=False, non_traced=None, max_cached_io=5, global_patching=True): +def start(name, save_path=".", verbose=False, dry_run=False, non_traced=None, + max_cached_io=5, global_patching=True): """Initialize and start LEAPP graph interpretation. ``name`` may be a bare graph name (``"my_graph"``), a relative path @@ -104,7 +106,8 @@ def stop(): manager.set_patches_applied(False) -def compile_graph(visualize=True, verbose=None, validate=True, dry_run=False, rtol=1e-3, atol=1e-5, strict=True): +def compile_graph(visualize=True, verbose=None, validate=True, dry_run=False, + rtol=1e-3, atol=1e-5, strict=True, graph_configs=None): """Compile and save the computational graph from traced nodes.""" manager = _MANAGER @@ -126,6 +129,12 @@ def compile_graph(visualize=True, verbose=None, validate=True, dry_run=False, rt graph = LeappGraph(manager.get_nodes(), manager.get_graph_name()) pipeline = graph.get_full_pipeline_description() + if graph_configs is not None: + if not isinstance(graph_configs, GraphConfigs): + raise TypeError( + "compile_graph(graph_configs=...) expects a GraphConfigs instance") + pipeline["pipeline"]["configs"] = graph_configs.to_dict() + initial_value_filename = None if not manager.is_dry_run(): initial_value_filename = graph.save_feedback_initial_values( diff --git a/leapp/utils/tensor_description.py b/leapp/utils/tensor_description.py index 3381831..7e696cf 100644 --- a/leapp/utils/tensor_description.py +++ b/leapp/utils/tensor_description.py @@ -16,7 +16,7 @@ # import collections.abc -from dataclasses import dataclass, fields +from dataclasses import dataclass, field, fields from typing import Optional, Any, Dict, Tuple, List import torch @@ -60,6 +60,45 @@ class CompactYamlDict(dict): lambda dumper, data: dumper.represent_mapping('tag:yaml.org,2002:map', data, flow_style=True)) +TEMPORAL_AXIS_SENTINEL = "__temporal_axis__" + + +@dataclass(frozen=True) +class TemporalAxis: + """Marks an element_names axis as temporal with a fixed period in ms.""" + + period_ms: float + + def __post_init__(self): + if self.period_ms <= 0: + raise ValueError("TemporalAxis period_ms must be positive") + + +@dataclass +class GraphConfigs: + """User-facing graph-level metadata for YAML serialization.""" + + frequency: Optional[float] = None + extra: Optional[Dict[str, Any]] = None + + def __post_init__(self): + if self.frequency is not None and self.frequency <= 0: + raise ValueError("GraphConfigs frequency must be positive when provided") + + def to_dict(self) -> Dict[str, Any]: + """Return non-None graph config fields with extra flattened.""" + result = {} + for f in fields(self): + if f.name == "extra": + continue + value = getattr(self, f.name) + if value is not None: + result[f.name] = value + if self.extra: + result.update(self.extra) + return result + + @dataclass class TensorSemantics: """User-facing semantic metadata for a tensor. @@ -82,16 +121,21 @@ class TensorSemantics: kind: Optional[InputKindEnum | OutputKindEnum | str] = None element_names: Optional[List] = None extra: Optional[Dict[str, Any]] = None + temporal_period_ms: Optional[float] = field(default=None, init=False) def __post_init__(self): '''error checking, auto conditioning''' + existing_period_ms = self.temporal_period_ms + self.temporal_period_ms = None if not is_tracable_tensor_type(self.ref): # this checks for both base types and traced types raise TypeError( f"TensorSemantics 'ref' must be a traceable tensor type " f"accepted types are: {TRACABLE_BASE_TYPES}" f"got {type(self.ref).__name__}") if self.element_names is not None: - self.element_names = self._normalize_element_names(self.element_names) + self.element_names, detected_period_ms = self._normalize_element_names( + self.element_names, allow_temporal_sentinel=existing_period_ms is not None) + self.temporal_period_ms = detected_period_ms or existing_period_ms def to_dict(self) -> Dict[str, Any]: @@ -117,7 +161,7 @@ def update(self, values: Dict[str, Any]): Known dataclass fields are set directly. Unknown keys are stored in the ``extra`` dict so they still appear in the serialized YAML. """ - valid_fields = {f.name for f in fields(self)} - self._INTERNAL_FIELDS + valid_fields = {f.name for f in fields(self) if f.init} - self._INTERNAL_FIELDS for key, value in values.items(): if key in valid_fields: setattr(self, key, value) @@ -128,24 +172,61 @@ def update(self, values: Dict[str, Any]): self.__post_init__() @staticmethod - def _normalize_element_names(element_names): - """Normalize element_names to CompactYamlList[CompactYamlList[str] | None]. - """ + def _normalize_element_names(element_names, allow_temporal_sentinel=False): + """Normalize element_names and extract temporal axis metadata.""" if isinstance(element_names, str): - return CompactYamlList([CompactYamlList([element_names])]) - elif isinstance(element_names, list): - if all(isinstance(item, str) for item in element_names): - return CompactYamlList([CompactYamlList(element_names)]) - elif all(isinstance(item, (list, type(None))) for item in element_names): - return CompactYamlList([ - CompactYamlList(item) if item is not None else None - for item in element_names - ]) + if element_names == TEMPORAL_AXIS_SENTINEL: + if allow_temporal_sentinel: + return CompactYamlList([element_names]), None + raise ValueError( + f"{TEMPORAL_AXIS_SENTINEL!r} is reserved for TemporalAxis") + return CompactYamlList([CompactYamlList([element_names])]), None + + if not isinstance(element_names, list): + return element_names, None + + temporal_period_ms = None + normalized = CompactYamlList() + has_axis_descriptors = False + has_temporal_axis = False + + for item in element_names: + if isinstance(item, TemporalAxis): + if has_temporal_axis: + raise ValueError("element_names can contain at most one TemporalAxis") + has_axis_descriptors = True + has_temporal_axis = True + temporal_period_ms = item.period_ms + normalized.append(TEMPORAL_AXIS_SENTINEL) + elif isinstance(item, str): + if item == TEMPORAL_AXIS_SENTINEL: + if not allow_temporal_sentinel: + raise ValueError( + f"{TEMPORAL_AXIS_SENTINEL!r} is reserved for TemporalAxis") + has_axis_descriptors = True + normalized.append(TEMPORAL_AXIS_SENTINEL) + else: + normalized.append(item) + elif isinstance(item, list): + if any(isinstance(child, TemporalAxis) for child in item): + raise ValueError("TemporalAxis must be an axis item, not nested in a list") + if any(child == TEMPORAL_AXIS_SENTINEL for child in item): + raise ValueError( + f"{TEMPORAL_AXIS_SENTINEL!r} must be a bare axis item, not nested in a list") + has_axis_descriptors = True + normalized.append(CompactYamlList(item)) + elif item is None: + has_axis_descriptors = True + normalized.append(None) else: _get_logger().warning( - "element_names has mixed types, expected List[List[str]]") - return element_names - return element_names + "element_names has mixed types, expected names or axis descriptors") + return element_names, temporal_period_ms + + if has_axis_descriptors: + return normalized, temporal_period_ms + + return CompactYamlList([CompactYamlList(normalized)]), None class TensorDescription: diff --git a/tests/functional_tests/test_config_generation.py b/tests/functional_tests/test_config_generation.py index 988d149..21b2f67 100644 --- a/tests/functional_tests/test_config_generation.py +++ b/tests/functional_tests/test_config_generation.py @@ -19,7 +19,7 @@ import yaml import torch import leapp -from leapp import TensorSemantics +from leapp import GraphConfigs, TensorSemantics, TemporalAxis from leapp.leapp import _MANAGER as annotate from leapp.utils.enums import InputKindEnum, OutputKindEnum from .base import LEAPPFunctionalTestBase @@ -358,6 +358,60 @@ def test_input_td_with_extra_fields_flattened_into_yaml(self): self.assertEqual(entry["frame"], "base") self.assertNotIn("extra", entry) + def test_graph_configs_appear_in_pipeline_yaml(self): + """Test GraphConfigs fields and extras are emitted as graph-level metadata.""" + tensor = torch.randn(1, 4) + + leapp.start(name=self.TEST_GRAPH_NAME) + traced = annotate.input_tensors("frequency_node", {"input": tensor}) + annotate.output_tensors("frequency_node", {"output": traced + 1.0}) + leapp.stop() + leapp.compile_graph( + visualize=False, + graph_configs=GraphConfigs( + frequency=50, + extra={"runtime": "isaac_lab"}, + ), + ) + + config = self._load_yaml() + + self.assertIn("pipeline", config) + self.assertEqual(config["pipeline"]["configs"]["frequency"], 50) + self.assertEqual(config["pipeline"]["configs"]["runtime"], "isaac_lab") + self.assertNotIn("frequency", config["pipeline"]) + self.assertNotIn("runtime", config["pipeline"]) + + def test_temporal_period_marker_appears_in_output_yaml(self): + """Test TemporalAxis emits temporal axis and period metadata.""" + tensor = torch.randn(2, 3) + names = ["hip", "knee", "ankle"] + + leapp.start(name=self.TEST_GRAPH_NAME) + traced = annotate.input_tensors("chunk_node", {"input": tensor}) + + annotate.output_tensors("chunk_node", [ + TensorSemantics( + name="actions", + ref=traced + 1.0, + element_names=[TemporalAxis(period_ms=100), names], + ), + TensorSemantics(name="plain_output", ref=traced - 1.0), + ]) + leapp.stop() + leapp.compile_graph(visualize=False) + + config = self._load_yaml() + _, outputs = self._get_node_io_from_yaml(config, "chunk_node") + + action_entry = self._find_io_by_name(outputs, "actions") + plain_entry = self._find_io_by_name(outputs, "plain_output") + + self.assertEqual(action_entry["element_names"], ["__temporal_axis__", names]) + self.assertEqual(action_entry["temporal_period_ms"], 100) + self.assertNotIn("temporal_period_ms", plain_entry) + self.assertNotIn("element_names", plain_entry) + # ========================================================================= # No metadata (baseline) # ========================================================================= @@ -387,6 +441,18 @@ def test_no_metadata_no_extra_yaml_fields(self): self.assertNotIn('source', output_entry) self.assertNotIn('element_names', output_entry) + def test_tensor_semantics_rejects_temporal_period_ms_argument(self): + """Test temporal period is only set through TemporalAxis.""" + tensor = torch.randn(2, 3) + + with self.assertRaises(TypeError): + TensorSemantics( + name="actions", + ref=tensor, + element_names=[["hip", "knee", "ankle"]], + temporal_period_ms=100, + ) + # ========================================================================= # Error cases # =========================================================================