wandb 0.19.12rc1__py3-none-win32.whl → 0.20.1__py3-none-win32.whl

wandb/sdk/internal/system/assets/trainium.py

@@ -126,8 +126,8 @@ class NeuronCoreStats:
                         self.raw_samples.append(raw_data)
                 process.kill()
                 process.wait()
-        except Exception as e:
-            logger.error("neuron-monitor failed: {}".format(e))
+        except Exception:
+            logger.exception("neuron-monitor failed")
 
     def __init__(
         self,
@@ -168,8 +168,8 @@ class NeuronCoreStats:
             self.shutdown_event.set()
             assert self.neuron_monitor_thread is not None
             self.neuron_monitor_thread.join()
-        except Exception as e:
-            logger.error("neuron-monitor thread failed to stop: {}".format(e))
+        except Exception:
+            logger.exception("neuron-monitor thread failed to stop")
         finally:
             self.neuron_monitor_thread = None
 
@@ -388,6 +388,6 @@ class Trainium:
                 pass
 
             return {self.name: neuron_hardware_info}
-        except Exception as e:
-            logger.error("neuron-monitor failed: {}".format(e))
+        except Exception:
+            logger.exception("neuron-monitor failed")
             return {}

wandb/sdk/internal/system/system_info.py

@@ -55,9 +55,7 @@ class SystemInfo:
         )
         program_absolute = os.path.join(root, program_relative)
         if not os.path.exists(program_absolute):
-            logger.warning(
-                "unable to save code -- can't find {}".format(program_absolute)
-            )
+            logger.warning(f"unable to save code -- can't find {program_absolute}")
             return None
         saved_program = os.path.join(self.settings.files_dir, "code", program_relative)
         self.saved_program = program_relative  # type: ignore
@@ -121,8 +119,8 @@ class SystemInfo:
             ValueError,
             subprocess.CalledProcessError,
             subprocess.TimeoutExpired,
-        ) as e:
-            logger.error("Error generating diff: {}".format(e))
+        ):
+            logger.exception("Error generating diff.")
         logger.debug("Saving git patches done")
 
     def _probe_git(self, data: Dict[str, Any]) -> Dict[str, Any]:
@@ -220,8 +218,8 @@ class SystemInfo:
                     stderr=subprocess.DEVNULL,
                     timeout=15,  # add timeout since conda env export could take a really long time
                 )
-        except Exception as e:
-            logger.exception(f"Error saving conda packages: {e}")
+        except Exception:
+            logger.exception("Error saving conda packages")
         logger.debug("Saving conda packages done")
 
     def publish(self, system_info: dict) -> None:

wandb/sdk/internal/system/system_monitor.py

@@ -177,8 +177,8 @@ class SystemMonitor:
             # publish telemetry
             self.publish_telemetry()
             self.aggregate_and_publish_asset_metrics()
-        except Exception as e:
-            logger.error(f"Error publishing last batch of metrics: {e}")
+        except Exception:
+            logger.exception("Error publishing last batch of metrics.")
 
     def start(self) -> None:
         self._shutdown_event.clear()
@@ -199,8 +199,8 @@ class SystemMonitor:
             asset.finish()
         try:
             self._process.join()
-        except Exception as e:
-            logger.error(f"Error joining system monitor process: {e}")
+        except Exception:
+            logger.exception("Error joining system monitor process.")
         self._process = None
 
     def probe(self, publish: bool = True) -> dict:

wandb/sdk/internal/tb_watcher.py

@@ -288,9 +288,9 @@ class TBDirWatcher:
     def _thread_except_body(self) -> None:
         try:
             self._thread_body()
-        except Exception as e:
+        except Exception:
             logger.exception("generic exception in TBDirWatcher thread")
-            raise e
+            raise
 
     def _thread_body(self) -> None:
         """Check for new events every second."""
@@ -394,9 +394,9 @@ class TBEventConsumer:
     def _thread_except_body(self) -> None:
         try:
             self._thread_body()
-        except Exception as e:
+        except Exception:
             logger.exception("generic exception in TBEventConsumer thread")
-            raise e
+            raise
 
     def _thread_body(self) -> None:
         while True:
@@ -490,9 +490,7 @@ class TBHistory:
                     dropped_keys.append(k)
                     del self._data[k]
             wandb.termwarn(
-                "Step {} exceeds max data limit, dropping {} of the largest keys:".format(
-                    self._step, len(dropped_keys)
-                )
+                f"Step {self._step} exceeds max data limit, dropping {len(dropped_keys)} of the largest keys:"
             )
             print("\t" + ("\n\t".join(dropped_keys)))  # noqa: T201
         self._data["_step"] = self._step

wandb/sdk/launch/_launch.py

@@ -55,7 +55,7 @@ def set_launch_logfile(logfile: str) -> None:
     _launch_logger.log(logging.INFO, "Internal agent logs printing to %s", logfile)
 
 
-def resolve_agent_config(  # noqa: C901
+def resolve_agent_config(
     entity: Optional[str],
     max_jobs: Optional[int],
     queues: Optional[Tuple[str]],

wandb/sdk/launch/_project_spec.py

@@ -7,13 +7,12 @@ import enum
 import json
 import logging
 import os
+import shlex
 import shutil
 import tempfile
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
 
-from six.moves import shlex_quote
-
 import wandb
 from wandb.apis.internal import Api
 from wandb.errors import CommError
@@ -121,7 +120,7 @@ class LaunchProject:
         ).get("base_image") or resource_args_build.get("cuda", {}).get("base_image")
         self.docker_image: Optional[str] = docker_config.get(
             "docker_image"
-        ) or launch_spec.get("image_uri")
+        ) or launch_spec.get("image_uri")  # type: ignore [assignment]
         self.docker_user_id = docker_config.get("user_id", 1000)
         self._entry_point: Optional[EntryPoint] = (
             None  # todo: keep multiple entrypoint support?
@@ -215,7 +214,7 @@ class LaunchProject:
             launch_spec.get("docker", {}),
             launch_spec.get("git", {}),
             launch_spec.get("overrides", {}),
-            launch_spec.get("resource", None),
+            launch_spec.get("resource", None),  # type: ignore [arg-type]
             launch_spec.get("resource_args", {}),
             launch_spec.get("run_id", None),
             launch_spec.get("sweep_id", {}),
@@ -487,30 +486,30 @@ class LaunchProject:
         return env_vars
 
     def parse_existing_requirements(self) -> str:
-        import pkg_resources
+        from packaging.requirements import InvalidRequirement, Requirement
 
         requirements_line = ""
         assert self.project_dir is not None
         base_requirements = os.path.join(self.project_dir, "requirements.txt")
         if os.path.exists(base_requirements):
             include_only = set()
-            with open(base_requirements) as f:
-                iter = pkg_resources.parse_requirements(f)
-                while True:
+            with open(base_requirements) as f2:
+                for line in f2:
+                    if line.strip() == "":
+                        continue
+
                     try:
-                        pkg = next(iter)
-                        if hasattr(pkg, "name"):
-                            name = pkg.name.lower()
-                        else:
-                            name = str(pkg)
-                        include_only.add(shlex_quote(name))
-                    except StopIteration:
-                        break
-                    # Different versions of pkg_resources throw different errors
-                    # just catch them all and ignore packages we can't parse
-                    except Exception as e:
-                        _logger.warn(f"Unable to parse requirements.txt: {e}")
+                        req = Requirement(line)
+                        name = req.name.lower()
+                        include_only.add(shlex.quote(name))
+                    except InvalidRequirement:
+                        _logger.warning(
+                            "Unable to parse line %s in requirements.txt",
+                            line,
+                            exc_info=True,
+                        )
                         continue
+
             requirements_line += "WANDB_ONLY_INCLUDE={} ".format(",".join(include_only))
             if "wandb" not in requirements_line:
                 wandb.termwarn(f"{LOG_PREFIX}wandb is not present in requirements.txt.")

wandb/sdk/launch/agent/agent.py

@@ -123,7 +123,7 @@ class InternalAgentLogger:
     def warn(self, message: str):
         if self._print_to_terminal:
             wandb.termwarn(f"{LOG_PREFIX}{message}")
-        _logger.warn(f"{LOG_PREFIX}{message}")
+        _logger.warning(f"{LOG_PREFIX}{message}")
 
     def info(self, message: str):
         if self._print_to_terminal:
@@ -589,7 +589,7 @@ class LaunchAgent:
                 )
                 if agent_response["stopPolling"]:
                     # shutdown process and all jobs if requested from ui
-                    raise KeyboardInterrupt
+                    raise KeyboardInterrupt  # noqa: TRY301
                 if self.num_running_jobs < self._max_jobs:
                     # only check for new jobs if we're not at max
                     job_and_queue = await self.get_job_and_queue()
@@ -850,7 +850,7 @@ class LaunchAgent:
                     )
                     return True
                 wandb.termlog(
-                    f"{LOG_PREFIX}Run {job_tracker.run_id} was preempted, requeueing..."
+                    f"{LOG_PREFIX}Run {job_tracker.run_id} was preempted, requeuing..."
                 )
 
                 if "sweep_id" in config:

wandb/sdk/launch/agent/config.py

@@ -7,7 +7,7 @@ from typing import List, Optional
 from pydantic import (  # type: ignore
     BaseModel,
     Field,
-    ValidationError,  # noqa: F401
+    ValidationError,
     root_validator,
     validator,
 )

wandb/sdk/launch/agent/job_status_tracker.py

@@ -48,6 +48,6 @@ class JobAndRunStatusTracker:
         check_stop = event_loop_thread_exec(api.api.check_stop_requested)
         try:
             return bool(await check_stop(self.project, self.entity, self.run_id))
-        except CommError as e:
-            _logger.error(f"CommError when checking if wandb run stopped: {e}")
+        except CommError:
+            _logger.exception("CommError when checking if wandb run stopped")
         return False

wandb/sdk/launch/builder/build.py

@@ -4,10 +4,9 @@ import logging
 import os
 import pathlib
 import shlex
+import shutil
 from typing import Any, Dict, List, Tuple
 
-from dockerpycreds.utils import find_executable  # type: ignore
-
 import wandb
 import wandb.env
 from wandb import docker
@@ -38,7 +37,7 @@ _WANDB_DOCKERFILE_NAME = "Dockerfile.wandb"
 
 async def validate_docker_installation() -> None:
     """Verify if Docker is installed on host machine."""
-    find_exec = event_loop_thread_exec(find_executable)
+    find_exec = event_loop_thread_exec(shutil.which)
     if not await find_exec("docker"):
         raise ExecutionError(
             "Could not find Docker executable. "

wandb/sdk/launch/builder/kaniko_builder.py

@@ -336,14 +336,15 @@ class KanikoBuilder(AbstractBuilder):
                 pod_name = get_pod_name_safe(k8s_job)
                 if pod_name:
                     msg += f" View logs with `kubectl logs -n {NAMESPACE} {pod_name}`."
-                raise Exception(msg)
+                raise Exception(msg)  # noqa: TRY301
             try:
                 pods_from_job = await core_v1.list_namespaced_pod(
                     namespace=NAMESPACE, label_selector=f"job-name={build_job_name}"
                 )
                 if len(pods_from_job.items) != 1:
-                    raise Exception(
-                        f"Expected 1 pod for job {build_job_name}, found {len(pods_from_job.items)}"
+                    raise Exception(  # noqa: TRY301
+                        f"Expected 1 pod for job {build_job_name},"
+                        f" found {len(pods_from_job.items)}"
                     )
                 pod_name = pods_from_job.items[0].metadata.name
                 logs = await core_v1.read_namespaced_pod_log(pod_name, NAMESPACE)
@@ -358,7 +359,7 @@ class KanikoBuilder(AbstractBuilder):
             wandb.termerror(
                 f"{LOG_PREFIX}Exception when creating Kubernetes resources: {e}\n"
             )
-            raise e
+            raise
         finally:
             wandb.termlog(f"{LOG_PREFIX}Cleaning up resources")
             try:

wandb/sdk/launch/environment/gcp_environment.py

@@ -94,8 +94,7 @@ class GcpEnvironment(AbstractEnvironment):
         region = config.get("region", None)
         if not region:
             raise LaunchError(
-                "Could not create GcpEnvironment from config. Missing 'region' "
-                "field."
+                "Could not create GcpEnvironment from config. Missing 'region' field."
             )
         return cls(region=region)
 

wandb/sdk/launch/registry/azure_container_registry.py

@@ -15,12 +15,12 @@ if TYPE_CHECKING:
     from azure.core.exceptions import ResourceNotFoundError  # type: ignore
 
 
-ContainerRegistryClient = get_module(  # noqa: F811
+ContainerRegistryClient = get_module(
     "azure.containerregistry",
     required="The azure-containerregistry package is required to use launch with Azure. Please install it with `pip install azure-containerregistry`.",
 ).ContainerRegistryClient
 
-ResourceNotFoundError = get_module(  # noqa: F811
+ResourceNotFoundError = get_module(
     "azure.core.exceptions",
     required="The azure-core package is required to use launch with Azure. Please install it with `pip install azure-core`.",
 ).ResourceNotFoundError

wandb/sdk/launch/registry/elastic_container_registry.py

@@ -14,11 +14,11 @@ from wandb.util import get_module
 
 _logger = logging.getLogger(__name__)
 
-botocore = get_module(  # noqa: F811
+botocore = get_module(
     "botocore",
     required="The boto3 package is required to use launch with AWS. Please install it with `pip install wandb[launch]`.",
 )
-boto3 = get_module(  # noqa: F811
+boto3 = get_module(
     "boto3",
     required="The boto3 package is required to use launch with AWS. Please install it with `pip install wandb[launch]`.",
 )

wandb/sdk/launch/registry/google_artifact_registry.py

@@ -17,16 +17,16 @@ from .abstract import AbstractRegistry
 
 _logger = logging.getLogger(__name__)
 
-google = get_module(  # noqa: F811
+google = get_module(
     "google",
     required="The google package is required to use launch with Google. Please install it with `pip install wandb[launch]`.",
 )
-google.auth = get_module(  # noqa: F811
+google.auth = get_module(
     "google.auth",
     required="The google-auth package is required to use launch with Google. Please install it with `pip install wandb[launch]`.",
 )
 
-google.cloud.artifactregistry = get_module(  # noqa: F811
+google.cloud.artifactregistry = get_module(
     "google.cloud.artifactregistry",
     required="The google-cloud-artifactregistry package is required to use launch with Google. Please install it with `pip install wandb[launch]`.",
 )

wandb/sdk/launch/runner/abstract.py

@@ -6,13 +6,12 @@ of runs launched in different environments (e.g. runs launched locally or in a c
 
 import logging
 import os
+import shutil
 import subprocess
 import sys
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Literal, Optional, Union
 
-from dockerpycreds.utils import find_executable  # type: ignore
-
 import wandb
 from wandb.apis.internal import Api
 from wandb.sdk.lib import runid
@@ -145,10 +144,11 @@ class AbstractRunner(ABC):
         self._namespace = runid.generate_id()
 
     def find_executable(
-        self, cmd: str
-    ) -> Any:  # should return a string, but mypy doesn't trust find_executable
+        self,
+        cmd: str,
+    ) -> Union[str, None]:
         """Cross platform utility for checking if a program is available."""
-        return find_executable(cmd)
+        return shutil.which(cmd)
 
     @property
     def api_key(self) -> Any:

wandb/sdk/launch/runner/kubernetes_monitor.py

@@ -5,10 +5,10 @@ import logging
 import traceback
 from typing import Any, Dict, List, Optional, Tuple, Union
 
-import kubernetes_asyncio  # type: ignore # noqa: F401
+import kubernetes_asyncio  # type: ignore
 import urllib3
 from kubernetes_asyncio import watch
-from kubernetes_asyncio.client import (  # type: ignore  # noqa: F401
+from kubernetes_asyncio.client import (  # type: ignore
     ApiException,
     BatchV1Api,
     CoreV1Api,

wandb/sdk/launch/runner/kubernetes_runner.py

@@ -491,7 +491,7 @@ class KubernetesRunner(AbstractRunner):
 
     async def run(
         self, launch_project: LaunchProject, image_uri: str
-    ) -> Optional[AbstractRun]:  # noqa: C901
+    ) -> Optional[AbstractRun]:
         """Execute a launch project on Kubernetes.
 
         Arguments:

wandb/sdk/launch/runner/sagemaker_runner.py

@@ -69,7 +69,7 @@ class SagemakerSubmittedRun(AbstractRun):
             )
             assert "events" in res
             return "\n".join(
-                [f'{event["timestamp"]}:{event["message"]}' for event in res["events"]]
+                [f"{event['timestamp']}:{event['message']}" for event in res["events"]]
             )
         except self.log_client.exceptions.ResourceNotFoundException:
             wandb.termwarn(
@@ -398,9 +398,7 @@ async def launch_sagemaker_job(
     wandb.termlog(
         f"{LOG_PREFIX}Run job submitted with arn: {resp.get('TrainingJobArn')}"
     )
-    url = "https://{region}.console.aws.amazon.com/sagemaker/home?region={region}#/jobs/{job_name}".format(
-        region=sagemaker_client.meta.region_name, job_name=training_job_name
-    )
+    url = f"https://{sagemaker_client.meta.region_name}.console.aws.amazon.com/sagemaker/home?region={sagemaker_client.meta.region_name}#/jobs/{training_job_name}"
     wandb.termlog(f"{LOG_PREFIX}See training job status at: {url}")
     return run
 

wandb/sdk/launch/runner/vertex_runner.py

@@ -49,12 +49,7 @@ class VertexSubmittedRun(AbstractRun):
         return self._job.project  # type: ignore
 
     def get_page_link(self) -> str:
-        return "{console_uri}/vertex-ai/locations/{region}/training/{job_id}?project={project}".format(
-            console_uri=GCP_CONSOLE_URI,
-            region=self.gcp_region,
-            job_id=self.id,
-            project=self.gcp_project,
-        )
+        return f"{GCP_CONSOLE_URI}/vertex-ai/locations/{self.gcp_region}/training/{self.id}?project={self.gcp_project}"
 
     async def wait(self) -> bool:
         # TODO: run this in a separate thread.
@@ -177,7 +172,7 @@ async def launch_vertex_job(
 ) -> VertexSubmittedRun:
     try:
         await environment.verify()
-        aiplatform = get_module(  # noqa: F811
+        aiplatform = get_module(
             "google.cloud.aiplatform",
             "VertexRunner requires google.cloud.aiplatform to be installed",
         )

wandb/sdk/launch/sweeps/__init__.py

@@ -27,7 +27,7 @@ def load_scheduler(scheduler_type: str) -> Any:
             f"{list(_WANDB_SCHEDULERS.keys())}, got: {scheduler_type}"
         )
 
-    log.warn(f"Loading dependencies for Scheduler of type: {scheduler_type}")
+    log.warning(f"Loading dependencies for Scheduler of type: {scheduler_type}")
     import_func = _WANDB_SCHEDULERS[scheduler_type]
     return import_func()
 

wandb/sdk/launch/sweeps/scheduler.py

@@ -354,7 +354,7 @@ class Scheduler(ABC):
             wandb.termlog(f"{LOG_PREFIX}Scheduler failed with exception {e}")
             self.state = SchedulerState.FAILED
             self.exit()
-            raise e
+            raise
         else:
             # scheduler succeeds if at runcap
             if self.state == SchedulerState.FLUSH_RUNS and self.at_runcap:
@@ -699,7 +699,7 @@ class Scheduler(ABC):
         if entry_point:
             wandb.termwarn(
                 f"{LOG_PREFIX}Sweep command {entry_point} will override"
-                f' {"job" if _job else "image_uri"} entrypoint'
+                f" {'job' if _job else 'image_uri'} entrypoint"
             )
 
         # override resource and args of job

wandb/sdk/launch/sweeps/utils.py

@@ -199,7 +199,7 @@ def create_sweep_command(command: Optional[List] = None) -> List:
             for m in matches[::-1]:
                 # Default to just leaving as is if environment variable does not exist
                 _var: str = os.environ.get(m.group(1), m.group(1))
-                command[i] = f"{command[i][:m.start()]}{_var}{command[i][m.end():]}"
+                command[i] = f"{command[i][: m.start()]}{_var}{command[i][m.end() :]}"
     return command
 
 
@@ -211,7 +211,7 @@ def create_sweep_command_args(command: Dict) -> Dict[str, Any]:
 
     """
     if "args" not in command:
-        raise ValueError('No "args" found in command: {}'.format(command))
+        raise ValueError(f'No "args" found in command: {command}')
     # four different formats of command args
     # (1) standard command line flags (e.g. --foo=bar)
     flags: List[str] = []
@@ -232,7 +232,7 @@ def create_sweep_command_args(command: Dict) -> Dict[str, Any]:
         try:
             _value: Any = config["value"]
         except KeyError:
-            raise ValueError('No "value" found for command["args"]["{}"]'.format(param))
+            raise ValueError(f'No "value" found for command["args"]["{param}"]')
 
         _flag: str = f"{param}={_value}"
         flags.append("--" + _flag)

wandb/sdk/launch/utils.py

@@ -323,8 +323,7 @@ def validate_launch_spec_source(launch_spec: Dict[str, Any]) -> None:
     docker_image = launch_spec.get("docker", {}).get("docker_image")
     if bool(job) == bool(docker_image):
         raise LaunchError(
-            "Exactly one of job or docker_image must be specified in the launch "
-            "spec."
+            "Exactly one of job or docker_image must be specified in the launch spec."
         )
 
 
@@ -628,9 +627,9 @@ def docker_image_exists(docker_image: str, should_raise: bool = False) -> bool:
     try:
         docker.run(["docker", "image", "inspect", docker_image])
         return True
-    except (docker.DockerError, ValueError) as e:
+    except (docker.DockerError, ValueError):
         if should_raise:
-            raise e
+            raise
         _logger.info("Base image not found. Generating new base image")
         return False
 

wandb/sdk/lib/apikey.py

@@ -21,6 +21,7 @@ import wandb
 from wandb.apis import InternalApi
 from wandb.errors import term
 from wandb.errors.links import url_registry
+from wandb.sdk import wandb_setup
 from wandb.util import _is_databricks, isatty, prompt_choices
 
 if TYPE_CHECKING:
@@ -237,9 +238,7 @@ def write_netrc(host: str, entity: str, key: str):
     _, key_suffix = key.split("-", 1) if "-" in key else ("", key)
     if len(key_suffix) != 40:
         raise ValueError(
-            "API-key must be exactly 40 characters long: {} ({} chars)".format(
-                key_suffix, len(key_suffix)
-            )
+            f"API-key must be exactly 40 characters long: {key_suffix} ({len(key_suffix)} chars)"
         )
 
     normalized_host = urlparse(host).netloc
@@ -276,7 +275,7 @@ def write_netrc(host: str, entity: str, key: str):
                     elif skip:
                         skip -= 1
                     else:
-                        f.write("{}\n".format(line))
+                        f.write(f"{line}\n")
 
             wandb.termlog(
                 f"Appending key for {normalized_host} to your netrc file: {netrc_path}"
@@ -311,16 +310,14 @@ def write_key(
     _, suffix = key.split("-", 1) if "-" in key else ("", key)
 
     if len(suffix) != 40:
-        raise ValueError(
-            "API key must be 40 characters long, yours was {}".format(len(key))
-        )
+        raise ValueError(f"API key must be 40 characters long, yours was {len(key)}")
 
     write_netrc(settings.base_url, "user", key)
 
 
 def api_key(settings: Settings | None = None) -> str | None:
     if settings is None:
-        settings = wandb.setup().settings
+        settings = wandb_setup.singleton().settings
     if settings.api_key:
         return settings.api_key
 

wandb/sdk/lib/config_util.py

@@ -66,13 +66,13 @@ def dict_from_config_file(
 ) -> Optional[Dict[str, Any]]:
     if not os.path.exists(filename):
         if must_exist:
-            raise ConfigError("config file {} doesn't exist".format(filename))
-        logger.debug("no default config file found in {}".format(filename))
+            raise ConfigError(f"config file {filename} doesn't exist")
+        logger.debug(f"no default config file found in {filename}")
         return None
     try:
         conf_file = open(filename)
     except OSError:
-        raise ConfigError("Couldn't read config file: {}".format(filename))
+        raise ConfigError(f"Couldn't read config file: {filename}")
     try:
         loaded = load_yaml(conf_file)
     except yaml.parser.ParserError:

wandb/sdk/lib/fsm.py

@@ -28,26 +28,11 @@ Usage:
     ```
 """
 
-import sys
 from abc import abstractmethod
 from dataclasses import dataclass
-from typing import (
-    Callable,
-    Dict,
-    Generic,
-    Optional,
-    Protocol,
-    Sequence,
-    Type,
-    TypeVar,
-    Union,
-    runtime_checkable,
-)
-
-if sys.version_info >= (3, 10):
-    from typing import TypeAlias
-else:
-    from typing_extensions import TypeAlias
+from typing import Callable, Dict, Generic, Optional, Sequence, Type, Union
+
+from typing_extensions import Protocol, TypeAlias, TypeVar, runtime_checkable
 
 T_FsmInputs = TypeVar("T_FsmInputs", contravariant=True)
 T_FsmContext = TypeVar("T_FsmContext")