incremental_table_util.py 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. from __future__ import annotations
  2. import time
  3. from typing import TYPE_CHECKING
  4. if TYPE_CHECKING:
  5. from wandb import Table
  6. from wandb.sdk.artifacts.artifact import Artifact
  7. from ..wandb_run import Run as LocalRun
  8. ART_TYPE = "wandb-run-incremental-table"
  9. def _get_artifact_name(run: LocalRun, key: str) -> str:
  10. from wandb.sdk.artifacts._internal_artifact import sanitize_artifact_name
  11. return sanitize_artifact_name(f"run-{run.id}-incr-{key}")
  12. def init_artifact(run: LocalRun, sanitized_key: str) -> Artifact:
  13. """Initialize a new artifact for an incremental table.
  14. Args:
  15. run: The wandb run associated with this artifact
  16. sanitized_key: Sanitized string key to identify the table
  17. Returns:
  18. A wandb Artifact configured for incremental table storage
  19. """
  20. from wandb.sdk.artifacts._internal_artifact import InternalArtifact
  21. artifact = InternalArtifact(
  22. _get_artifact_name(run, sanitized_key),
  23. ART_TYPE,
  24. incremental=True,
  25. )
  26. return artifact
  27. def get_entry_name(incr_table: Table, key: str) -> str:
  28. """Generate a unique entry name for a table increment.
  29. Args:
  30. run: The wandb run associated with this table
  31. incr_table: The incremental table being updated
  32. key: String key for the table entry
  33. Returns:
  34. A unique string name for the table entry
  35. """
  36. epoch = time.time_ns() // 1_000_000
  37. return f"{incr_table._increment_num}-{epoch}.{key}"