Skip to content

Commit c4b9ec2

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
feat: Update Experiment run names to be compatible with One Platform API
PiperOrigin-RevId: 689075065
1 parent 5cc9052 commit c4b9ec2

File tree

2 files changed

+52
-1
lines changed

2 files changed

+52
-1
lines changed

google/cloud/aiplatform/tensorboard/uploader_utils.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -43,14 +43,15 @@
4343
TensorboardServiceClient = tensorboard_service_client.TensorboardServiceClient
4444
DEFAULT_RUN_NAME = "default"
4545
DEFAULT_PROFILE_RUN_NAME = "profile"
46+
REPLACE_PATTERN = "[^a-zA-Z0-9\n-]"
4647

4748
logger = tb_logging.get_logger()
4849
logger.setLevel(logging.WARNING)
4950

5051

5152
def reformat_run_name(run_name: str) -> str:
5253
"""Reformats the run name to be compatible with One Platform."""
53-
return run_name.replace("/", "-").replace("_", "-")
54+
return re.sub(REPLACE_PATTERN, "-", run_name)
5455

5556

5657
class RateLimiter:

tests/unit/aiplatform/test_uploader.py

+50
Original file line numberDiff line numberDiff line change
@@ -674,6 +674,56 @@ def test_start_uploading_scalars(
674674
self.assertEqual(mock_tracker.tensors_tracker.call_count, 0)
675675
self.assertEqual(mock_tracker.blob_tracker.call_count, 0)
676676

677+
@parameterized.parameters(
678+
{"run_name": "test-run-1", "one_platform_run_name": "test-run-1"},
679+
{"run_name": "test/run/2", "one_platform_run_name": "test-run-2"},
680+
{"run_name": "test.run.3", "one_platform_run_name": "test-run-3"},
681+
{"run_name": "test_run_4", "one_platform_run_name": "test-run-4"},
682+
{"run_name": "test/.run_5", "one_platform_run_name": "test--run-5"},
683+
)
684+
@patch.object(
685+
uploader_utils.OnePlatformResourceManager,
686+
"get_run_resource_name",
687+
autospec=True,
688+
)
689+
@patch.object(metadata, "_experiment_tracker", autospec=True)
690+
@patch.object(experiment_resources, "Experiment", autospec=True)
691+
def test_start_uploading_scalars_run_names_updated_to_one_platform_names(
692+
self,
693+
experiment_resources_mock,
694+
experiment_tracker_mock,
695+
run_resource_mock,
696+
run_name,
697+
one_platform_run_name,
698+
):
699+
experiment_resources_mock.get.return_value = _TEST_EXPERIMENT_NAME
700+
experiment_tracker_mock.set_experiment.return_value = _TEST_EXPERIMENT_NAME
701+
experiment_tracker_mock.set_tensorboard.return_value = (
702+
_TEST_TENSORBOARD_RESOURCE_NAME
703+
)
704+
run_resource_mock.return_value = _TEST_ONE_PLATFORM_RUN_NAME
705+
uploader = _create_uploader(
706+
logdir=_TEST_LOG_DIR_NAME,
707+
)
708+
uploader.create_experiment()
709+
710+
mock_logdir_loader = mock.create_autospec(logdir_loader.LogdirLoader)
711+
mock_logdir_loader.get_run_events.side_effect = [
712+
{
713+
run_name: _apply_compat(
714+
[_scalar_event("1.1", 5.0), _scalar_event("1.2", 5.0)]
715+
),
716+
},
717+
AbortUploadError,
718+
]
719+
720+
with mock.patch.object(
721+
uploader, "_logdir_loader", mock_logdir_loader
722+
), self.assertRaises(AbortUploadError):
723+
uploader.start_uploading()
724+
experiment_runs = uploader._experiment_runs
725+
self.assertIn(one_platform_run_name, experiment_runs)
726+
677727
@parameterized.parameters(
678728
{
679729
"existing_experiment": None,

0 commit comments

Comments
 (0)