wandb 0.16.6__py3-none-any.whl → 0.17.0__py3-none-any.whl

wandb/sdk/launch/runner/sagemaker_runner.py

@@ -1,4 +1,5 @@
 """Implementation of the SageMakerRunner class."""
+
 import asyncio
 import logging
 from typing import Any, Dict, List, Optional, cast
@@ -11,8 +12,7 @@ from wandb.apis.internal import Api
 from wandb.sdk.launch.environment.aws_environment import AwsEnvironment
 from wandb.sdk.launch.errors import LaunchError
 
-from .._project_spec import EntryPoint, LaunchProject, get_entry_point_command
-from ..builder.build import get_env_vars_dict
+from .._project_spec import EntryPoint, LaunchProject
 from ..registry.abstract import AbstractRegistry
 from ..utils import (
     LOG_PREFIX,
@@ -67,6 +67,7 @@ class SagemakerSubmittedRun(AbstractRun):
                 logGroupName="/aws/sagemaker/TrainingJobs",
                 logStreamName=log_name,
             )
+            assert "events" in res
             return "\n".join(
                 [f'{event["timestamp"]}:{event["message"]}' for event in res["events"]]
             )
@@ -220,12 +221,12 @@ class SageMakerRunner(AbstractRunner):
         launch_project.fill_macros(image_uri)
         _logger.info("Connecting to sagemaker client")
         entry_point = (
-            launch_project.override_entrypoint
-            or launch_project.get_single_entry_point()
-        )
-        command_args = get_entry_point_command(
-            entry_point, launch_project.override_args
+            launch_project.override_entrypoint or launch_project.get_job_entry_point()
         )
+        command_args = []
+        if entry_point is not None:
+            command_args += entry_point.command
+        command_args += launch_project.override_args
         if command_args:
             command_str = " ".join(command_args)
             wandb.termlog(
@@ -324,16 +325,16 @@ def build_sagemaker_args(
     sagemaker_args["TrainingJobName"] = training_job_name
     entry_cmd = entry_point.command if entry_point else []
 
-    sagemaker_args[
-        "AlgorithmSpecification"
-    ] = merge_image_uri_with_algorithm_specification(
-        given_sagemaker_args.get(
-            "AlgorithmSpecification",
-            given_sagemaker_args.get("algorithm_specification"),
-        ),
-        image_uri,
-        entry_cmd,
-        args,
+    sagemaker_args["AlgorithmSpecification"] = (
+        merge_image_uri_with_algorithm_specification(
+            given_sagemaker_args.get(
+                "AlgorithmSpecification",
+                given_sagemaker_args.get("algorithm_specification"),
+            ),
+            image_uri,
+            entry_cmd,
+            args,
+        )
     )
 
     sagemaker_args["RoleArn"] = role_arn
@@ -348,18 +349,18 @@ def build_sagemaker_args(
 
     if sagemaker_args.get("ResourceConfig") is None:
         raise LaunchError(
-            "Sagemaker launcher requires a ResourceConfig Sagemaker resource argument"
+            "Sagemaker launcher requires a ResourceConfig resource argument"
         )
 
     if sagemaker_args.get("StoppingCondition") is None:
         raise LaunchError(
-            "Sagemaker launcher requires a StoppingCondition Sagemaker resource argument"
+            "Sagemaker launcher requires a StoppingCondition resource argument"
         )
 
     given_env = given_sagemaker_args.get(
         "Environment", sagemaker_args.get("environment", {})
     )
-    calced_env = get_env_vars_dict(launch_project, api, max_env_length)
+    calced_env = launch_project.get_env_vars_dict(api, max_env_length)
     total_env = {**calced_env, **given_env}
     sagemaker_args["Environment"] = total_env
 

wandb/sdk/launch/runner/vertex_runner.py

@@ -8,8 +8,7 @@ if False:
 from wandb.apis.internal import Api
 from wandb.util import get_module
 
-from .._project_spec import LaunchProject, get_entry_point_command
-from ..builder.build import get_env_vars_dict
+from .._project_spec import LaunchProject
 from ..environment.gcp_environment import GcpEnvironment
 from ..errors import LaunchError
 from ..registry.abstract import AbstractRegistry
@@ -113,14 +112,16 @@ class VertexRunner(AbstractRunner):
         synchronous: bool = self.backend_config[PROJECT_SYNCHRONOUS]
 
         entry_point = (
-            launch_project.override_entrypoint
-            or launch_project.get_single_entry_point()
+            launch_project.override_entrypoint or launch_project.get_job_entry_point()
         )
 
         # TODO: Set entrypoint in each container
-        entry_cmd = get_entry_point_command(entry_point, launch_project.override_args)
-        env_vars = get_env_vars_dict(
-            launch_project=launch_project,
+        entry_cmd = []
+        if entry_point is not None:
+            entry_cmd += entry_point.command
+        entry_cmd += launch_project.override_args
+
+        env_vars = launch_project.get_env_vars_dict(
             api=self._api,
             max_env_length=MAX_ENV_LENGTHS[self.__class__.__name__],
         )

wandb/sdk/launch/sweeps/scheduler.py

@@ -1,4 +1,5 @@
 """Abstract Scheduler class."""
+
 import asyncio
 import base64
 import copy
@@ -407,7 +408,7 @@ class Scheduler(ABC):
         return count
 
     def _try_load_executable(self) -> bool:
-        """Check existance of valid executable for a run.
+        """Check existence of valid executable for a run.
 
         logs and returns False when job is unreachable
         """
@@ -422,7 +423,7 @@ class Scheduler(ABC):
                 return False
             return True
         elif self._kwargs.get("image_uri"):
-            # TODO(gst): check docker existance? Use registry in launch config?
+            # TODO(gst): check docker existence? Use registry in launch config?
             return True
         else:
             return False
@@ -610,7 +611,7 @@ class Scheduler(ABC):
                     f"Failed to get runstate for run ({run_id}). Error: {traceback.format_exc()}"
                 )
                 run_state = RunState.FAILED
-            else:  # first time we get unknwon state
+            else:  # first time we get unknown state
                 run_state = RunState.UNKNOWN
         except (AttributeError, ValueError):
             wandb.termwarn(

wandb/sdk/launch/sweeps/scheduler_sweep.py

@@ -1,4 +1,5 @@
 """Scheduler for classic wandb Sweeps."""
+
 import logging
 from pprint import pformat as pf
 from typing import Any, Dict, List, Optional
@@ -58,7 +59,7 @@ class SweepScheduler(Scheduler):
         return None
 
     def _get_sweep_commands(self, worker_id: int) -> List[Dict[str, Any]]:
-        """Helper to recieve sweep command from backend."""
+        """Helper to receive sweep command from backend."""
         # AgentHeartbeat wants a Dict of runs which are running or queued
         _run_states: Dict[str, bool] = {}
         for run_id, run in self._yield_runs():

wandb/sdk/launch/sweeps/utils.py

@@ -217,7 +217,7 @@ def create_sweep_command_args(command: Dict) -> Dict[str, Any]:
     flags: List[str] = []
     # (2) flags without hyphens (e.g. foo=bar)
     flags_no_hyphens: List[str] = []
-    # (3) flags with false booleans ommited  (e.g. --foo)
+    # (3) flags with false booleans omitted  (e.g. --foo)
     flags_no_booleans: List[str] = []
     # (4) flags as a dictionary (used for constructing a json)
     flags_dict: Dict[str, Any] = {}
@@ -257,7 +257,7 @@ def make_launch_sweep_entrypoint(
     """Use args dict from create_sweep_command_args to construct entrypoint.
 
     If replace is True, remove macros from entrypoint, fill them in with args
-    and then return the args in seperate return value.
+    and then return the args in separate return value.
     """
     if not command:
         return None, None
@@ -296,7 +296,7 @@ def check_job_exists(public_api: "PublicApi", job: Optional[str]) -> bool:
 
 
 def get_previous_args(
-    run_spec: Dict[str, Any]
+    run_spec: Dict[str, Any],
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     """Parse through previous scheduler run_spec.
 

wandb/sdk/launch/utils.py

@@ -1,4 +1,3 @@
-# heavily inspired by https://github.com/mlflow/mlflow/blob/master/mlflow/projects/utils.py
 import asyncio
 import json
 import logging
@@ -16,7 +15,6 @@ import wandb
 import wandb.docker as docker
 from wandb import util
 from wandb.apis.internal import Api
-from wandb.errors import CommError
 from wandb.sdk.launch.errors import LaunchError
 from wandb.sdk.launch.git_reference import GitReference
 from wandb.sdk.launch.wandb_reference import WandbReference
@@ -32,7 +30,6 @@ FAILED_PACKAGES_REGEX = re.compile(
 )
 
 if TYPE_CHECKING:  # pragma: no cover
-    from wandb.sdk.artifacts.artifact import Artifact
     from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker
 
 
@@ -57,15 +54,15 @@ API_KEY_REGEX = r"WANDB_API_KEY=\w+(-\w+)?"
 MACRO_REGEX = re.compile(r"\$\{(\w+)\}")
 
 AZURE_CONTAINER_REGISTRY_URI_REGEX = re.compile(
-    r"(?:https://)?([\w]+)\.azurecr\.io/([\w\-]+):?(.*)"
+    r"^(?:https://)?([\w]+)\.azurecr\.io/(?P<repository>[\w\-]+):?(?P<tag>.*)"
 )
 
 ELASTIC_CONTAINER_REGISTRY_URI_REGEX = re.compile(
-    r"^(?P<account>.*)\.dkr\.ecr\.(?P<region>.*)\.amazonaws\.com/(?P<repository>.*)/?$"
+    r"^(?:https://)?(?P<account>[\w-]+)\.dkr\.ecr\.(?P<region>[\w-]+)\.amazonaws\.com/(?P<repository>[\w-]+):?(?P<tag>.*)$"
 )
 
 GCP_ARTIFACT_REGISTRY_URI_REGEX = re.compile(
-    r"^(?P<region>[\w-]+)-docker\.pkg\.dev/(?P<project>[\w-]+)/(?P<repository>[\w-]+)/(?P<image_name>[\w-]+)$",
+    r"^(?:https://)?(?P<region>[\w-]+)-docker\.pkg\.dev/(?P<project>[\w-]+)/(?P<repository>[\w-]+)/?(?P<image_name>[\w-]+)?(?P<tag>:.*)?$",
     re.IGNORECASE,
 )
 
@@ -316,16 +313,13 @@ def construct_launch_spec(
 
 
 def validate_launch_spec_source(launch_spec: Dict[str, Any]) -> None:
-    uri = launch_spec.get("uri")
     job = launch_spec.get("job")
     docker_image = launch_spec.get("docker", {}).get("docker_image")
-
-    if not bool(uri) and not bool(job) and not bool(docker_image):
-        raise LaunchError("Must specify a uri, job or docker image")
-    elif bool(uri) and bool(docker_image):
-        raise LaunchError("Found both uri and docker-image, only one can be set")
-    elif sum(map(bool, [uri, job, docker_image])) > 1:
-        raise LaunchError("Must specify exactly one of uri, job or image")
+    if bool(job) == bool(docker_image):
+        raise LaunchError(
+            "Exactly one of job or docker_image must be specified in the launch "
+            "spec."
+        )
 
 
 def parse_wandb_uri(uri: str) -> Tuple[str, str, str]:
@@ -336,77 +330,6 @@ def parse_wandb_uri(uri: str) -> Tuple[str, str, str]:
     return (ref.entity, ref.project, ref.run_id)
 
 
-def is_bare_wandb_uri(uri: str) -> bool:
-    """Check that a wandb uri is valid.
-
-    URI must be in the format
-    `/<entity>/<project>/runs/<run_name>[other stuff]`
-    or
-    `/<entity>/<project>/artifacts/job/<job_name>[other stuff]`.
-    """
-    _logger.info(f"Checking if uri {uri} is bare...")
-    return uri.startswith("/") and WandbReference.is_uri_job_or_run(uri)
-
-
-def fetch_wandb_project_run_info(
-    entity: str, project: str, run_name: str, api: Api
-) -> Any:
-    _logger.info("Fetching run info...")
-    try:
-        result = api.get_run_info(entity, project, run_name)
-    except CommError:
-        result = None
-    if result is None:
-        raise LaunchError(
-            f"Run info is invalid or doesn't exist for {api.settings('base_url')}/{entity}/{project}/runs/{run_name}"
-        )
-    if result.get("codePath") is None:
-        # TODO: we don't currently expose codePath in the runInfo endpoint, this downloads
-        # it from wandb-metadata.json if we can.
-        metadata = api.download_url(
-            project, "wandb-metadata.json", run=run_name, entity=entity
-        )
-        if metadata is not None:
-            _, response = api.download_file(metadata["url"])
-            data = response.json()
-            result["codePath"] = data.get("codePath")
-            result["cudaVersion"] = data.get("cuda", None)
-
-    return result
-
-
-def download_entry_point(
-    entity: str, project: str, run_name: str, api: Api, entry_point: str, dir: str
-) -> bool:
-    metadata = api.download_url(
-        project, f"code/{entry_point}", run=run_name, entity=entity
-    )
-    if metadata is not None:
-        _, response = api.download_file(metadata["url"])
-        with util.fsync_open(os.path.join(dir, entry_point), "wb") as file:
-            for data in response.iter_content(chunk_size=1024):
-                file.write(data)
-        return True
-    return False
-
-
-def download_wandb_python_deps(
-    entity: str, project: str, run_name: str, api: Api, dir: str
-) -> Optional[str]:
-    reqs = api.download_url(project, "requirements.txt", run=run_name, entity=entity)
-    if reqs is not None:
-        _logger.info("Downloading python dependencies")
-        _, response = api.download_file(reqs["url"])
-
-        with util.fsync_open(
-            os.path.join(dir, "requirements.frozen.txt"), "wb"
-        ) as file:
-            for data in response.iter_content(chunk_size=1024):
-                file.write(data)
-        return "requirements.frozen.txt"
-    return None
-
-
 def get_local_python_deps(
     dir: str, filename: str = "requirements.local.txt"
 ) -> Optional[str]:
@@ -498,19 +421,6 @@ def validate_wandb_python_deps(
     _logger.warning("Unable to validate local python dependencies")
 
 
-def fetch_project_diff(
-    entity: str, project: str, run_name: str, api: Api
-) -> Optional[str]:
-    """Fetches project diff from wandb servers."""
-    _logger.info("Searching for diff.patch")
-    patch = None
-    try:
-        (_, _, patch, _) = api.run_config(project, run_name, entity)
-    except CommError:
-        pass
-    return patch
-
-
 def apply_patch(patch_string: str, dst_dir: str) -> None:
     """Applies a patch file to a directory."""
     _logger.info("Applying diff.patch")
@@ -531,17 +441,6 @@ def apply_patch(patch_string: str, dst_dir: str) -> None:
         raise wandb.Error("Failed to apply diff.patch associated with run.")
 
 
-def _make_refspec_from_version(version: Optional[str]) -> List[str]:
-    """Create a refspec that checks for the existence of origin/main and the version."""
-    if version:
-        return [f"+{version}"]
-
-    return [
-        "+refs/heads/main*:refs/remotes/origin/main*",
-        "+refs/heads/master*:refs/remotes/origin/master*",
-    ]
-
-
 def _fetch_git_repo(dst_dir: str, uri: str, version: Optional[str]) -> Optional[str]:
     """Clones the git repo at ``uri`` into ``dst_dir``.
 
@@ -561,13 +460,6 @@ def _fetch_git_repo(dst_dir: str, uri: str, version: Optional[str]) -> Optional[
     return version
 
 
-def merge_parameters(
-    higher_priority_params: Dict[str, Any], lower_priority_params: Dict[str, Any]
-) -> Dict[str, Any]:
-    """Merge the contents of two dicts, keeping values from higher_priority_params if there are conflicts."""
-    return {**lower_priority_params, **higher_priority_params}
-
-
 def convert_jupyter_notebook_to_script(fname: str, project_dir: str) -> str:
     nbconvert = wandb.util.get_module(
         "nbconvert", "nbformat and nbconvert are required to use launch with notebooks"
@@ -597,25 +489,6 @@ def convert_jupyter_notebook_to_script(fname: str, project_dir: str) -> str:
     return new_name
 
 
-def check_and_download_code_artifacts(
-    entity: str, project: str, run_name: str, internal_api: Api, project_dir: str
-) -> Optional["Artifact"]:
-    _logger.info("Checking for code artifacts")
-    public_api = wandb.PublicApi(
-        overrides={"base_url": internal_api.settings("base_url")}
-    )
-
-    run = public_api.run(f"{entity}/{project}/{run_name}")
-    run_artifacts = run.logged_artifacts()
-
-    for artifact in run_artifacts:
-        if hasattr(artifact, "type") and artifact.type == "code":
-            artifact.download(project_dir)
-            return artifact  # type: ignore
-
-    return None
-
-
 def to_camel_case(maybe_snake_str: str) -> str:
     if "_" not in maybe_snake_str:
         return maybe_snake_str
@@ -623,11 +496,6 @@ def to_camel_case(maybe_snake_str: str) -> str:
     return "".join(x.title() if x else "_" for x in components)
 
 
-def run_shell(args: List[str]) -> Tuple[str, str]:
-    out = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    return out.stdout.decode("utf-8").strip(), out.stderr.decode("utf-8").strip()
-
-
 def validate_build_and_registry_configs(
     build_config: Dict[str, Any], registry_config: Dict[str, Any]
 ) -> None:
@@ -864,3 +732,10 @@ def get_entrypoint_file(entrypoint: List[str]) -> Optional[str]:
     if len(entrypoint) < 2:
         return None
     return entrypoint[1]
+
+
+def get_current_python_version() -> Tuple[str, str]:
+    full_version = sys.version.split()[0].split(".")
+    major = full_version[0]
+    version = ".".join(full_version[:2]) if len(full_version) >= 2 else major + ".0"
+    return version, major

wandb/sdk/lib/_settings_toposort_generated.py

@@ -13,7 +13,6 @@ else:
 _Setting = Literal[
     "_args",
     "_aws_lambda",
-    "_async_upload_concurrency_limit",
     "_cli_only_mode",
     "_code_path_local",
     "_colab",
@@ -25,7 +24,6 @@ _Setting = Literal[
     "_disable_update_check",
     "_disable_viewer",
     "_disable_machine_info",
-    "_except_exit",
     "_executable",
     "_extra_http_headers",
     "_file_stream_retry_max",
@@ -126,7 +124,6 @@ _Setting = Literal[
     "login_timeout",
     "mode",
     "notebook_name",
-    "problem",
     "program",
     "program_abspath",
     "program_relpath",
@@ -179,7 +176,6 @@ _Setting = Literal[
 ]
 
 SETTINGS_TOPOLOGICALLY_SORTED: Final[Tuple[_Setting, ...]] = (
-    "_async_upload_concurrency_limit",
     "_service_wait",
     "_stats_sample_rate_seconds",
     "_stats_samples_to_average",
@@ -189,7 +185,6 @@ SETTINGS_TOPOLOGICALLY_SORTED: Final[Tuple[_Setting, ...]] = (
     "console",
     "job_source",
     "mode",
-    "problem",
     "project",
     "run_id",
     "start_method",

wandb/sdk/lib/fsm.py

@@ -52,43 +52,39 @@ T_FsmContext_contra = TypeVar("T_FsmContext_contra", contravariant=True)
 @runtime_checkable
 class FsmStateCheck(Protocol[T_FsmInputs]):
     @abstractmethod
-    def on_check(self, inputs: T_FsmInputs) -> None:
-        ...  # pragma: no cover
+    def on_check(self, inputs: T_FsmInputs) -> None: ...  # pragma: no cover
 
 
 @runtime_checkable
 class FsmStateOutput(Protocol[T_FsmInputs]):
     @abstractmethod
-    def on_state(self, inputs: T_FsmInputs) -> None:
-        ...  # pragma: no cover
+    def on_state(self, inputs: T_FsmInputs) -> None: ...  # pragma: no cover
 
 
 @runtime_checkable
 class FsmStateEnter(Protocol[T_FsmInputs]):
     @abstractmethod
-    def on_enter(self, inputs: T_FsmInputs) -> None:
-        ...  # pragma: no cover
+    def on_enter(self, inputs: T_FsmInputs) -> None: ...  # pragma: no cover
 
 
 @runtime_checkable
 class FsmStateEnterWithContext(Protocol[T_FsmInputs, T_FsmContext_contra]):
     @abstractmethod
-    def on_enter(self, inputs: T_FsmInputs, context: T_FsmContext_contra) -> None:
-        ...  # pragma: no cover
+    def on_enter(
+        self, inputs: T_FsmInputs, context: T_FsmContext_contra
+    ) -> None: ...  # pragma: no cover
 
 
 @runtime_checkable
 class FsmStateStay(Protocol[T_FsmInputs]):
     @abstractmethod
-    def on_stay(self, inputs: T_FsmInputs) -> None:
-        ...  # pragma: no cover
+    def on_stay(self, inputs: T_FsmInputs) -> None: ...  # pragma: no cover
 
 
 @runtime_checkable
 class FsmStateExit(Protocol[T_FsmInputs, T_FsmContext_cov]):
     @abstractmethod
-    def on_exit(self, inputs: T_FsmInputs) -> T_FsmContext_cov:
-        ...  # pragma: no cover
+    def on_exit(self, inputs: T_FsmInputs) -> T_FsmContext_cov: ...  # pragma: no cover
 
 
 # It would be nice if python provided optional protocol members, but it doesnt as described here:

wandb/sdk/lib/gitlib.py

@@ -14,7 +14,7 @@ try:
         Repo,
     )
 except ImportError:
-    Repo = None
+    Repo = None  # type: ignore
 
 if TYPE_CHECKING:
     from git import Repo
@@ -121,7 +121,7 @@ class GitRepo:
         # TODO: Saw a user getting a Unicode decode error when parsing refs,
         # more details on implementing a real fix in [WB-4064]
         try:
-            if len(self.repo.refs) > 0:
+            if len(self.repo.refs) > 0:  # type: ignore[arg-type]
                 return self.repo.head.commit.hexsha
             else:
                 return self.repo.git.show_ref("--head").split(" ")[0]
@@ -140,7 +140,7 @@ class GitRepo:
         if not self.repo:
             return None
         try:
-            return self.repo.remotes[self.remote_name]
+            return self.repo.remotes[self.remote_name]  # type: ignore[index]
         except IndexError:
             return None
 
@@ -200,7 +200,7 @@ class GitRepo:
                     possible_relatives.append(tracking_branch.commit)
 
             if not possible_relatives:
-                for branch in self.repo.branches:
+                for branch in self.repo.branches:  # type: ignore[attr-defined]
                     tracking_branch = branch.tracking_branch()
                     if tracking_branch is not None:
                         possible_relatives.append(tracking_branch.commit)

wandb/sdk/lib/import_hooks.py

@@ -143,7 +143,7 @@ class _ImportHookChainedLoader:
         # None, so handle None as well. The module may not support attribute
         # assignment, in which case we simply skip it. Note that we also deal
         # with __loader__ not existing at all. This is to future proof things
-        # due to proposal to remove the attribue as described in the GitHub
+        # due to proposal to remove the attribute as described in the GitHub
         # issue at https://github.com/python/cpython/issues/77458. Also prior
         # to Python 3.3, the __loader__ attribute was only set if a custom
         # module loader was used. It isn't clear whether the attribute still

wandb/sdk/lib/lazyloader.py

@@ -1,6 +1,5 @@
 """module lazyloader."""
 
-
 import importlib
 import sys
 import types

wandb/sdk/lib/proto_util.py

@@ -12,7 +12,28 @@ if TYPE_CHECKING:  # pragma: no cover
 
 
 def dict_from_proto_list(obj_list: "RepeatedCompositeFieldContainer") -> Dict[str, Any]:
-    return {item.key: json.loads(item.value_json) for item in obj_list}
+    result: Dict[str, Any] = {}
+
+    for item in obj_list:
+        # Start from the root of the result dict
+        current_level = result
+
+        if len(item.nested_key) > 0:
+            keys = list(item.nested_key)
+        else:
+            keys = [item.key]
+
+        for key in keys[:-1]:
+            if key not in current_level:
+                current_level[key] = {}
+            # Move the reference deeper into the nested dictionary
+            current_level = current_level[key]
+
+        # Set the value at the final key location, parsing JSON from the value_json field
+        final_key = keys[-1]
+        current_level[final_key] = json.loads(item.value_json)
+
+    return result
 
 
 def _result_from_record(record: "pb.Record") -> "pb.Result":
@@ -29,7 +50,7 @@ def _assign_end_offset(record: "pb.Record", end_offset: int) -> None:
 
 
 def proto_encode_to_dict(
-    pb_obj: Union["tpb.TelemetryRecord", "pb.MetricRecord"]
+    pb_obj: Union["tpb.TelemetryRecord", "pb.MetricRecord"],
 ) -> Dict[int, Any]:
     data: Dict[int, Any] = dict()
     fields = pb_obj.ListFields()