xmanager-slurm 0.3.1__py3-none-any.whl → 0.4.0__py3-none-any.whl

xm_slurm/__init__.py

@@ -4,6 +4,7 @@ from xm_slurm.executables import Dockerfile, DockerImage
 from xm_slurm.executors import Slurm, SlurmSpec
 from xm_slurm.experiment import (
     Artifact,
+    SlurmExperiment,
     create_experiment,
     get_current_experiment,
     get_current_work_unit,
@@ -14,8 +15,8 @@ from xm_slurm.packageables import (
     docker_container,
     docker_image,
     mamba_container,
-    pdm_container,
     python_container,
+    uv_container,
 )
 from xm_slurm.resources import JobRequirements, ResourceQuantity, ResourceType
 
@@ -35,10 +36,11 @@ __all__ = [
     "get_experiment",
     "JobRequirements",
     "mamba_container",
-    "pdm_container",
+    "uv_container",
     "python_container",
     "ResourceQuantity",
     "ResourceType",
     "Slurm",
     "SlurmSpec",
+    "SlurmExperiment",
 ]

xm_slurm/api.py

@@ -1,15 +1,27 @@
 import dataclasses
+import enum
 import functools
 import importlib.util
 import logging
 import os
-import time
 import typing
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from pathlib import Path
 from typing import Any
 
+from sqlalchemy import Column, ForeignKey, Integer, String, create_engine
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import relationship, sessionmaker
+
 logger = logging.getLogger(__name__)
 
 
+class ExperimentUnitRole(enum.Enum):
+    WORK_UNIT = enum.auto()
+    AUX_UNIT = enum.auto()
+
+
 @dataclasses.dataclass(kw_only=True, frozen=True)
 class ExperimentPatchModel:
     title: str | None = None
@@ -22,7 +34,7 @@ class ExperimentPatchModel:
 class SlurmJobModel:
     name: str
     slurm_job_id: int
-    slurm_cluster: str
+    slurm_ssh_config: str
 
 
 @dataclasses.dataclass(kw_only=True, frozen=True)
@@ -32,15 +44,21 @@ class ArtifactModel:
 
 
 @dataclasses.dataclass(kw_only=True, frozen=True)
-class WorkUnitPatchModel:
-    wid: int
+class ExperimentUnitModel:
+    identity: str
+    args: str | None = None
+    jobs: list[SlurmJobModel] = dataclasses.field(default_factory=list)
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class ExperimentUnitPatchModel:
     identity: str | None
     args: str | None = None
 
 
 @dataclasses.dataclass(kw_only=True, frozen=True)
-class WorkUnitModel(WorkUnitPatchModel):
-    jobs: list[SlurmJobModel] = dataclasses.field(default_factory=list)
+class WorkUnitModel(ExperimentUnitModel):
+    wid: int
     artifacts: list[ArtifactModel] = dataclasses.field(default_factory=list)
 
 
@@ -55,49 +73,54 @@ class ExperimentModel:
     artifacts: list[ArtifactModel]
 
 
-class XManagerAPI:
+class XManagerAPI(ABC):
+    @abstractmethod
     def get_experiment(self, xid: int) -> ExperimentModel:
-        del xid
-        raise NotImplementedError("`get_experiment` is not implemented without a storage backend.")
+        pass
 
+    @abstractmethod
     def delete_experiment(self, experiment_id: int) -> None:
-        del experiment_id
-        logger.debug("`delete_experiment` is not implemented without a storage backend.")
+        pass
 
+    @abstractmethod
     def insert_experiment(self, experiment: ExperimentPatchModel) -> int:
-        del experiment
-        logger.debug("`insert_experiment` is not implemented without a storage backend.")
-        return int(time.time() * 10**3)
+        pass
 
+    @abstractmethod
     def update_experiment(self, experiment_id: int, experiment_patch: ExperimentPatchModel) -> None:
-        del experiment_id, experiment_patch
-        logger.debug("`update_experiment` is not implemented without a storage backend.")
+        pass
 
+    @abstractmethod
     def insert_job(self, experiment_id: int, work_unit_id: int, job: SlurmJobModel) -> None:
-        del experiment_id, work_unit_id, job
-        logger.debug("`insert_job` is not implemented without a storage backend.")
+        pass
 
-    def insert_work_unit(self, experiment_id: int, work_unit: WorkUnitPatchModel) -> None:
-        del experiment_id, work_unit
-        logger.debug("`insert_work_unit` is not implemented without a storage backend.")
+    @abstractmethod
+    def insert_work_unit(self, experiment_id: int, work_unit: WorkUnitModel) -> None:
+        pass
 
+    @abstractmethod
+    def update_work_unit(
+        self, experiment_id: int, work_unit_id: int, patch: ExperimentUnitPatchModel
+    ) -> None:
+        pass
+
+    @abstractmethod
     def delete_work_unit_artifact(self, experiment_id: int, work_unit_id: int, name: str) -> None:
-        del experiment_id, work_unit_id, name
-        logger.debug("`delete_work_unit_artifact` is not implemented without a storage backend.")
+        pass
 
+    @abstractmethod
     def insert_work_unit_artifact(
         self, experiment_id: int, work_unit_id: int, artifact: ArtifactModel
     ) -> None:
-        del experiment_id, work_unit_id, artifact
-        logger.debug("`insert_work_unit_artifact` is not implemented without a storage backend.")
+        pass
 
+    @abstractmethod
     def delete_experiment_artifact(self, experiment_id: int, name: str) -> None:
-        del experiment_id, name
-        logger.debug("`delete_experiment_artifact` is not implemented without a storage backend.")
+        pass
 
+    @abstractmethod
     def insert_experiment_artifact(self, experiment_id: int, artifact: ArtifactModel) -> None:
-        del experiment_id, artifact
-        logger.debug("`insert_experiment_artifact` is not implemented without a storage backend.")
+        pass
 
 
 class XManagerWebAPI(XManagerAPI):
@@ -180,8 +203,6 @@ class XManagerWebAPI(XManagerAPI):
             update_experiment as _update_experiment,
         )
 
-        m = self.models.ExperimentPatch(**dataclasses.asdict(experiment_patch))
-
         _update_experiment.sync(
             experiment_id,
             client=self.client,
@@ -198,7 +219,7 @@ class XManagerWebAPI(XManagerAPI):
             body=self.models.SlurmJob(**dataclasses.asdict(job)),
         )
 
-    def insert_work_unit(self, experiment_id: int, work_unit: WorkUnitPatchModel) -> None:
+    def insert_work_unit(self, experiment_id: int, work_unit: WorkUnitModel) -> None:
         from xm_slurm_api_client.api.work_unit import (  # type: ignore
             insert_work_unit as _insert_work_unit,
         )
@@ -244,6 +265,253 @@ class XManagerWebAPI(XManagerAPI):
         )
 
 
+Base = declarative_base()
+
+
+class Experiment(Base):
+    __tablename__ = "experiments"
+
+    id = Column(Integer, primary_key=True)
+    title = Column(String)
+    description = Column(String)
+    note = Column(String)
+    tags = Column(String)
+    work_units = relationship("WorkUnit", back_populates="experiment")
+    artifacts = relationship("Artifact", back_populates="experiment")
+
+
+class WorkUnit(Base):
+    __tablename__ = "work_units"
+
+    id = Column(Integer, primary_key=True)
+    experiment_id = Column(Integer, ForeignKey("experiments.id"))
+    wid = Column(Integer)
+    identity = Column(String)
+    args = Column(String)
+    experiment = relationship("Experiment", back_populates="work_units")
+    jobs = relationship("SlurmJob", back_populates="work_unit")
+    artifacts = relationship("Artifact", back_populates="work_unit")
+
+
+class SlurmJob(Base):
+    __tablename__ = "slurm_jobs"
+
+    id = Column(Integer, primary_key=True)
+    work_unit_id = Column(Integer, ForeignKey("work_units.id"))
+    name = Column(String)
+    slurm_job_id = Column(Integer)
+    slurm_ssh_config = Column(String)
+    work_unit = relationship("WorkUnit", back_populates="jobs")
+
+
+class Artifact(Base):
+    __tablename__ = "artifacts"
+
+    id = Column(Integer, primary_key=True)
+    experiment_id = Column(Integer, ForeignKey("experiments.id"))
+    work_unit_id = Column(Integer, ForeignKey("work_units.id"))
+    name = Column(String)
+    uri = Column(String)
+    experiment = relationship("Experiment", back_populates="artifacts")
+    work_unit = relationship("WorkUnit", back_populates="artifacts")
+
+
+class XManagerSqliteAPI(XManagerAPI):
+    def __init__(self):
+        if "XM_SLURM_STATE_DIR" in os.environ:
+            db_path = Path(os.environ["XM_SLURM_STATE_DIR"]) / "db.sqlite3"
+        else:
+            db_path = Path.home() / ".local" / "state" / "xm-slurm" / "db.sqlite3"
+        logging.debug("Looking for db at: ", db_path)
+        db_path.parent.mkdir(parents=True, exist_ok=True)
+        engine = create_engine(f"sqlite:///{db_path}")
+        Base.metadata.create_all(engine)
+        self.Session = sessionmaker(bind=engine)
+
+    @contextmanager
+    def session_scope(self):
+        session = self.Session()
+        try:
+            yield session
+            session.commit()
+        except:
+            session.rollback()
+            raise
+        finally:
+            session.close()
+
+    def get_experiment(self, xid: int) -> ExperimentModel:
+        with self.session_scope() as session:
+            experiment = session.query(Experiment).filter(Experiment.id == xid).first()
+            if not experiment:
+                raise ValueError(f"Experiment with id {xid} not found")
+
+            work_units = []
+            for wu in experiment.work_units:
+                jobs = [
+                    SlurmJobModel(
+                        name=job.name,
+                        slurm_job_id=job.slurm_job_id,
+                        slurm_ssh_config=job.slurm_ssh_config,
+                    )
+                    for job in wu.jobs
+                ]
+                artifacts = [
+                    ArtifactModel(name=artifact.name, uri=artifact.uri) for artifact in wu.artifacts
+                ]
+                work_units.append(
+                    WorkUnitModel(
+                        wid=wu.wid,
+                        identity=wu.identity,
+                        args=wu.args,
+                        jobs=jobs,
+                        artifacts=artifacts,
+                    )
+                )
+
+            artifacts = [
+                ArtifactModel(name=artifact.name, uri=artifact.uri)
+                for artifact in experiment.artifacts
+            ]
+
+            return ExperimentModel(
+                title=experiment.title,
+                description=experiment.description,
+                note=experiment.note,
+                tags=experiment.tags.split(",") if experiment.tags else None,
+                work_units=work_units,
+                artifacts=artifacts,
+            )
+
+    def delete_experiment(self, experiment_id: int) -> None:
+        with self.session_scope() as session:
+            experiment = session.query(Experiment).filter(Experiment.id == experiment_id).first()
+            if experiment:
+                session.delete(experiment)
+
+    def insert_experiment(self, experiment: ExperimentPatchModel) -> int:
+        with self.session_scope() as session:
+            new_experiment = Experiment(
+                title=experiment.title,
+                description=experiment.description,
+                note=experiment.note,
+                tags=",".join(experiment.tags) if experiment.tags else None,
+            )
+            session.add(new_experiment)
+            session.flush()
+            return new_experiment.id
+
+    def update_experiment(self, experiment_id: int, experiment_patch: ExperimentPatchModel) -> None:
+        with self.session_scope() as session:
+            experiment = session.query(Experiment).filter(Experiment.id == experiment_id).first()
+            if experiment:
+                if experiment_patch.title is not None:
+                    experiment.title = experiment_patch.title
+                if experiment_patch.description is not None:
+                    experiment.description = experiment_patch.description
+                if experiment_patch.note is not None:
+                    experiment.note = experiment_patch.note
+                if experiment_patch.tags is not None:
+                    experiment.tags = ",".join(experiment_patch.tags)
+
+    def insert_job(self, experiment_id: int, work_unit_id: int, job: SlurmJobModel) -> None:
+        with self.session_scope() as session:
+            work_unit = (
+                session.query(WorkUnit)
+                .filter_by(experiment_id=experiment_id, wid=work_unit_id)
+                .first()
+            )
+            if work_unit:
+                new_job = SlurmJob(
+                    work_unit_id=work_unit.id,
+                    name=job.name,
+                    slurm_job_id=job.slurm_job_id,
+                    slurm_ssh_config=job.slurm_ssh_config,
+                )
+                session.add(new_job)
+            else:
+                raise ValueError(
+                    f"Work unit with id {work_unit_id} not found in experiment {experiment_id}"
+                )
+
+    def insert_work_unit(self, experiment_id: int, work_unit: WorkUnitModel) -> None:
+        with self.session_scope() as session:
+            new_work_unit = WorkUnit(
+                experiment_id=experiment_id,
+                wid=work_unit.wid,
+                identity=work_unit.identity,
+                args=work_unit.args,
+            )
+            session.add(new_work_unit)
+            for job in work_unit.jobs:
+                new_job = SlurmJob(
+                    work_unit_id=new_work_unit.id,
+                    name=job.name,
+                    slurm_job_id=job.slurm_job_id,
+                    slurm_ssh_config=job.slurm_ssh_config,
+                )
+                session.add(new_job)
+            for artifact in work_unit.artifacts:
+                new_artifact = Artifact(
+                    work_unit_id=new_work_unit.id, name=artifact.name, uri=artifact.uri
+                )
+                session.add(new_artifact)
+
+    def update_work_unit(
+        self, experiment_id: int, work_unit_id: int, patch: ExperimentUnitPatchModel
+    ) -> None:
+        with self.session_scope() as session:
+            work_unit = (
+                session.query(WorkUnit)
+                .filter(WorkUnit.experiment_id == experiment_id, WorkUnit.wid == work_unit_id)
+                .first()
+            )
+
+            if work_unit:
+                if patch.identity is not None:
+                    work_unit.identity = patch.identity
+                if patch.args is not None:
+                    work_unit.args = patch.args
+            else:
+                raise ValueError(
+                    f"Work unit with id {work_unit_id} not found in experiment {experiment_id}"
+                )
+
+    def delete_work_unit_artifact(self, experiment_id: int, work_unit_id: int, name: str) -> None:
+        with self.session_scope() as session:
+            artifact = (
+                session.query(Artifact)
+                .filter(Artifact.work_unit_id == work_unit_id, Artifact.name == name)
+                .first()
+            )
+            if artifact:
+                session.delete(artifact)
+
+    def insert_work_unit_artifact(
+        self, experiment_id: int, work_unit_id: int, artifact: ArtifactModel
+    ) -> None:
+        with self.session_scope() as session:
+            new_artifact = Artifact(work_unit_id=work_unit_id, name=artifact.name, uri=artifact.uri)
+            session.add(new_artifact)
+
+    def delete_experiment_artifact(self, experiment_id: int, name: str) -> None:
+        with self.session_scope() as session:
+            artifact = (
+                session.query(Artifact)
+                .filter(Artifact.experiment_id == experiment_id, Artifact.name == name)
+                .first()
+            )
+            if artifact:
+                session.delete(artifact)
+
+    def insert_experiment_artifact(self, experiment_id: int, artifact: ArtifactModel) -> None:
+        with self.session_scope() as session:
+            new_artifact = Artifact(
+                experiment_id=experiment_id, name=artifact.name, uri=artifact.uri
+            )
+            session.add(new_artifact)
+
+
 @functools.cache
 def client() -> XManagerAPI:
     if importlib.util.find_spec("xm_slurm_api_client") is not None:
@@ -257,5 +525,4 @@ def client() -> XManagerAPI:
                 "Disabling XManager API client."
             )
 
-    logger.debug("xm_slurm_api_client not found... skipping logging to the API.")
-    return XManagerAPI()
+    return XManagerSqliteAPI()

xm_slurm/batching.py

@@ -21,7 +21,7 @@ def stack_bound_arguments(
     signature: inspect.Signature, bound_arguments: Sequence[inspect.BoundArguments]
 ) -> inspect.BoundArguments:
     """Stacks bound arguments into a single bound arguments object."""
-    stacked_args = collections.OrderedDict()
+    stacked_args = collections.OrderedDict[str, Any]()
     for bound_args in bound_arguments:
         for name, value in bound_args.arguments.items():
             stacked_args.setdefault(name, [])
@@ -59,7 +59,7 @@ class batch(Generic[R]):
         self.loop: asyncio.AbstractEventLoop | None = None
         self.process_batch_task: asyncio.Task | None = None
 
-        self.queue = asyncio.Queue()
+        self.queue = asyncio.Queue[Request]()
 
     async def _process_batch(self):
         assert self.loop is not None
@@ -128,12 +128,12 @@ class batch(Generic[R]):
         # until then this is just a hack
         return asyncio.coroutines._is_coroutine  # type: ignore
 
-    async def __call__(self, *args, **kwargs) -> asyncio.Future[R]:
+    async def __call__(self, *args, **kwargs) -> R:
         if self.loop is None and self.process_batch_task is None:
             self.loop = asyncio.get_event_loop()
             self.process_batch_task = self.loop.create_task(self._process_batch())
 
-        future = asyncio.Future()
+        future = asyncio.Future[R]()
         bound_args = self.signature.bind(*args, **kwargs)
         self.queue.put_nowait(Request(args=bound_args, future=future))
         return await future

xm_slurm/config.py

@@ -2,12 +2,15 @@ import dataclasses
 import enum
 import functools
 import getpass
+import json
 import os
 import pathlib
 from typing import Literal, Mapping, NamedTuple
 
 import asyncssh
 
+from xm_slurm import constants
+
 
 class ContainerRuntime(enum.Enum):
     """The container engine to use."""
@@ -46,15 +49,103 @@ class PublicKey(NamedTuple):
     key: str
 
 
-@dataclasses.dataclass(frozen=True, kw_only=True)
-class SlurmClusterConfig:
-    name: str
-
+@dataclasses.dataclass
+class SlurmSSHConfig:
     host: str
     host_public_key: PublicKey | None = None
     user: str | None = None
     port: int | None = None
 
+    @functools.cached_property
+    def known_hosts(self) -> asyncssh.SSHKnownHosts | None:
+        if self.host_public_key is None:
+            return None
+
+        return asyncssh.import_known_hosts(
+            f"[{self.host}]:{self.port} {self.host_public_key.algorithm} {self.host_public_key.key}"
+        )
+
+    @functools.cached_property
+    def config(self) -> asyncssh.config.SSHConfig:
+        ssh_config_paths = []
+        if (ssh_config := pathlib.Path.home() / ".ssh" / "config").exists():
+            ssh_config_paths.append(ssh_config)
+        if (xm_ssh_config_var := os.environ.get("XM_SLURM_SSH_CONFIG")) and (
+            xm_ssh_config := pathlib.Path(xm_ssh_config_var).expanduser()
+        ).exists():
+            ssh_config_paths.append(xm_ssh_config)
+
+        config = asyncssh.config.SSHClientConfig.load(
+            None,
+            ssh_config_paths,
+            False,
+            getpass.getuser(),
+            self.user or (),
+            self.host,
+            self.port or (),
+        )
+        if config.get("Hostname") is None and (
+            constants.DOMAIN_NAME_REGEX.match(self.host)
+            or constants.IPV4_REGEX.match(self.host)
+            or constants.IPV6_REGEX.match(self.host)
+        ):
+            config._options["Hostname"] = self.host
+        elif config.get("Hostname") is None:
+            raise RuntimeError(
+                f"Failed to parse hostname from host `{self.host}` using "
+                f"SSH configs: {', '.join(map(str, ssh_config_paths))} and "
+                f"provided hostname `{self.host}` isn't a valid domain name "
+                "or IPv{4,6} address."
+            )
+
+        if config.get("User") is None:
+            raise RuntimeError(
+                f"We could not find a user for the cluster configuration: `{self.host}`. "
+                "No user was specified in the configuration and we could not parse "
+                f"any users for host `{config.get('Hostname')}` from the SSH configs: "
+                f"{', '.join(map(lambda h: f'`{h}`', ssh_config_paths))}. Please either specify a user "
+                "in the configuration or add a user to your SSH configuration under the block "
+                f"`Host {config.get('Hostname')}`."
+            )
+
+        return config
+
+    @functools.cached_property
+    def connection_options(self) -> asyncssh.SSHClientConnectionOptions:
+        options = asyncssh.SSHClientConnectionOptions(config=None)
+        options.prepare(last_config=self.config, known_hosts=self.known_hosts)
+        return options
+
+    def serialize(self):
+        return json.dumps({
+            "host": self.host,
+            "host_public_key": self.host_public_key,
+            "user": self.user,
+            "port": self.port,
+        })
+
+    @classmethod
+    def deserialize(cls, data):
+        data = json.loads(data)
+        return cls(
+            host=data["host"],
+            host_public_key=PublicKey(*data["host_public_key"])
+            if data["host_public_key"]
+            else None,
+            user=data["user"],
+            port=data["port"],
+        )
+
+    def __hash__(self):
+        return hash((self.host, self.host_public_key, self.user, self.port))
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class SlurmClusterConfig:
+    name: str
+
+    ssh: SlurmSSHConfig
+
     # Job submission directory
     cwd: str | None = None
 
@@ -81,7 +172,9 @@ class SlurmClusterConfig:
     )
 
     # Resource mapping
-    resources: Mapping[str, "xm_slurm.ResourceType"] = dataclasses.field(default_factory=dict)  # type: ignore # noqa: F821
+    resources: Mapping["xm_slurm.ResourceType", str] = dataclasses.field(default_factory=dict)  # type: ignore # noqa: F821
+
+    features: Mapping["xm_slurm.FeatureType", str] = dataclasses.field(default_factory=dict)  # type: ignore # noqa: F821
 
     def __post_init__(self) -> None:
         for src, dst in self.mounts.items():
@@ -99,57 +192,9 @@ class SlurmClusterConfig:
             if not pathlib.Path(dst).is_absolute():
                 raise ValueError(f"Mount destination must be an absolute path: {dst}")
 
-    @functools.cached_property
-    def ssh_known_hosts(self) -> asyncssh.SSHKnownHosts | None:
-        if self.host_public_key is None:
-            return None
-
-        return asyncssh.import_known_hosts(
-            f"[{self.host}]:{self.port} {self.host_public_key.algorithm} {self.host_public_key.key}"
-        )
-
-    @functools.cached_property
-    def ssh_config(self) -> asyncssh.config.SSHConfig:
-        ssh_config_paths = []
-        if (ssh_config := pathlib.Path.home() / ".ssh" / "config").exists():
-            ssh_config_paths.append(ssh_config)
-        if (xm_ssh_config := os.environ.get("XM_SLURM_SSH_CONFIG")) and (
-            xm_ssh_config := pathlib.Path(xm_ssh_config).expanduser()
-        ).exists():
-            ssh_config_paths.append(xm_ssh_config)
-
-        config = asyncssh.config.SSHClientConfig.load(
-            None,
-            ssh_config_paths,
-            True,
-            getpass.getuser(),
-            self.user or (),
-            self.host or (),
-            self.port or (),
-        )
-
-        if config.get("Hostname") is None:
-            raise RuntimeError(
-                f"Failed to parse hostname from host `{self.host}` using SSH configs: {', '.join(map(str, ssh_config_paths))}"
-            )
-        if config.get("User") is None:
-            raise RuntimeError(
-                f"Failed to parse user from SSH configs: {', '.join(map(str, ssh_config_paths))}"
-            )
-
-        return config
-
-    @functools.cached_property
-    def ssh_connection_options(self) -> asyncssh.SSHClientConnectionOptions:
-        options = asyncssh.SSHClientConnectionOptions(config=None)
-        options.prepare(last_config=self.ssh_config, known_hosts=self.ssh_known_hosts)
-        return options
-
     def __hash__(self):
         return hash((
-            self.host,
-            self.user,
-            self.port,
+            self.ssh,
             self.cwd,
             self.prolog,
             self.epilog,

xm_slurm/constants.py

@@ -0,0 +1,15 @@
+import re
+
+IMAGE_URI_REGEX = re.compile(
+    r"^(?P<scheme>(?:[^:]+://)?)?(?P<domain>[^/]+)(?P<path>/[^:]*)?(?::(?P<tag>[^@]+))?@?(?P<digest>.+)?$"
+)
+
+DOMAIN_NAME_REGEX = re.compile(
+    r"^(?!-)(?!.*--)[A-Za-z0-9-]{1,63}(?<!-)(\.[A-Za-z0-9-]{1,63})*(\.[A-Za-z]{2,})$"
+)
+IPV4_REGEX = re.compile(
+    r"^((25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$"
+)
+IPV6_REGEX = re.compile(
+    r"^(([0-9a-fA-F]{1,4}:){7}([0-9a-fA-F]{1,4}|:)|(([0-9a-fA-F]{1,4}:){1,7}|:):(([0-9a-fA-F]{1,4}:){1,6}|:)|(([0-9a-fA-F]{1,4}:){1,6}|:):(([0-9a-fA-F]{1,4}:){1,5}|:)|(([0-9a-fA-F]{1,4}:){1,5}|:):(([0-9a-fA-F]{1,4}:){1,4}|:)|(([0-9a-fA-F]{1,4}:){1,4}|:):(([0-9a-fA-F]{1,4}:){1,3}|:)|(([0-9a-fA-F]{1,4}:){1,3}|:):(([0-9a-fA-F]{1,4}:){1,2}|:)|(([0-9a-fA-F]{1,4}:){1,2}|:):([0-9a-fA-F]{1,4}|:)|([0-9a-fA-F]{1,4}|:):([0-9a-fA-F]{1,4}|:)(:([0-9a-fA-F]{1,4}|:)){1,6})$"
+)