[Jobs] Add network groups to hf jobs run (#4833)

* [Jobs] Add network groups to hf jobs run

* Trim network groups guide sections

---------

Co-authored-by: Lucain <lucain@huggingface.co>

Author: co42

docs/source/en/guides/cli.md

@@ -2116,6 +2116,20 @@ Pass `--ssh` to `hf jobs run` (or `hf jobs uv run`) to make the Job's container
 
 Only users with write access to the Job's namespace are allowed in (the Job creator, or members of the owner organization), authenticated by an SSH public key registered at https://huggingface.co/settings/keys.
 
+### Network groups
+
+Pass `--network-group <name>` to `hf jobs run` (or `hf jobs uv run`) to let Jobs of the same owner reach each other on every port. Inside each member, `$HF_NETWORK_GROUP_HOSTNAME` resolves to every Job in the group, and `${HF_NETWORK_GROUP_PREFIX}<alias>` to the members that claimed an alias with `--network-alias <alias>`:
+
+```bash
+# Start a server, reachable by the other members of the group as "master"
+>>> hf jobs run --detach --network-group train --network-alias master python:3.12 python -m http.server 8000
+
+# Start a client in the same group
+>>> hf jobs run --detach --network-group train python:3.12 sh -c 'curl --retry 10 --retry-connrefused "http://${HF_NETWORK_GROUP_PREFIX}master:8000/"'
+```
+
+Members are resolvable before they are ready, so connect with retries.
+
 ### UV Scripts (Experimental)
 
 Run UV scripts (Python scripts with inline dependencies) on HF infrastructure. UV scripts are Python scripts that include their dependencies directly in the file using a special comment syntax.

docs/source/en/guides/jobs.md

@@ -372,6 +372,27 @@ Connect from a terminal with `hf jobs ssh <job_id>` (or directly with `ssh <job_
 
 Only users with write access to the Job's namespace are allowed in (the Job creator, or members of the owner organization), authenticated by an SSH public key registered at https://huggingface.co/settings/keys.
 
+## Network groups
+
+Pass `network_group="<name>"` to [`run_job`] (or [`run_uv_job`]) to let Jobs of the same owner reach each other on every port. Inside each member, `HF_NETWORK_GROUP_HOSTNAME` resolves to every Job in the group, and `${HF_NETWORK_GROUP_PREFIX}<alias>` to the members that claimed an alias with `network_aliases=[...]`:
+
+```python
+>>> from huggingface_hub import run_job
+>>> server = run_job(
+...     image="python:3.12",
+...     command=["python", "-m", "http.server", "8000"],
+...     network_group="train",
+...     network_aliases=["master"],
+... )
+>>> client = run_job(
+...     image="python:3.12",
+...     command=["sh", "-c", 'curl --retry 10 --retry-connrefused "http://${HF_NETWORK_GROUP_PREFIX}master:8000/"'],
+...     network_group="train",
+... )
+```
+
+Members are resolvable before they are ready, so connect with retries.
+
 ## Configure Job Timeout
 
 Jobs have a default timeout (30 minutes), after which they will automatically stop. This is important to know when running long-running tasks like model training.

docs/source/en/package_reference/cli.md

@@ -2505,6 +2505,8 @@ $ hf jobs run [OPTIONS] IMAGE COMMAND...
 * `-d, --detach`: Run the Job in the background and print the Job ID.
 * `--expose INTEGER`: Expose a container port through the jobs proxy. Repeat the flag for multiple ports (e.g. `--expose 8000 --expose 8001`). Each exposed port is reachable on the public jobs domain; access requires an HF token with read access to the job's namespace.
 * `--ssh`: Make the job's container reachable over SSH. Connect with `hf jobs ssh <job_id>`. Requires an SSH public key registered on https://huggingface.co/settings/keys.
+* `--network-group TEXT`: Join a network group. Jobs of the same owner sharing a group are placed together and reach each other on every port. Inside each member, `$HF_NETWORK_GROUP_HOSTNAME` resolves to every member. Lowercase alphanumerics and dashes, 46 characters max.
+* `--network-alias TEXT`: Claim an alias in the network group. Members reach the jobs claiming it at `${HF_NETWORK_GROUP_PREFIX}<alias>`. Repeat the flag for several aliases. Requires `--network-group`.
 * `--resource-group-id TEXT`: The ID of the resource group to create the Job in. Used to control access to resources within an organization and for cost attribution/spending-limit features.
 * `--namespace TEXT`: The namespace where the job will be running. Defaults to the current user's namespace.
 * `--token TEXT`: A User Access Token generated from https://huggingface.co/settings/tokens.
@@ -2977,6 +2979,8 @@ $ hf jobs uv run [OPTIONS] SCRIPT [SCRIPT_ARGS]...
 * `-d, --detach`: Run the Job in the background and print the Job ID.
 * `--expose INTEGER`: Expose a container port through the jobs proxy. Repeat the flag for multiple ports (e.g. `--expose 8000 --expose 8001`). Each exposed port is reachable on the public jobs domain; access requires an HF token with read access to the job's namespace.
 * `--ssh`: Make the job's container reachable over SSH. Connect with `hf jobs ssh <job_id>`. Requires an SSH public key registered on https://huggingface.co/settings/keys.
+* `--network-group TEXT`: Join a network group. Jobs of the same owner sharing a group are placed together and reach each other on every port. Inside each member, `$HF_NETWORK_GROUP_HOSTNAME` resolves to every member. Lowercase alphanumerics and dashes, 46 characters max.
+* `--network-alias TEXT`: Claim an alias in the network group. Members reach the jobs claiming it at `${HF_NETWORK_GROUP_PREFIX}<alias>`. Repeat the flag for several aliases. Requires `--network-group`.
 * `--resource-group-id TEXT`: The ID of the resource group to create the Job in. Used to control access to resources within an organization and for cost attribution/spending-limit features.
 * `--namespace TEXT`: The namespace where the job will be running. Defaults to the current user's namespace.
 * `--token TEXT`: A User Access Token generated from https://huggingface.co/settings/tokens.

src/huggingface_hub/_jobs_api.py

@@ -581,8 +581,12 @@ def _create_job_spec(
     volumes: list[Volume] | None = None,
     expose: list[int] | None = None,
     ssh: bool = False,
+    network_group: str | None = None,
+    network_aliases: list[str] | None = None,
     resource_group_id: str | None = None,
 ) -> dict[str, Any]:
+    if network_aliases and not network_group:
+        raise ValueError("`network_aliases` requires `network_group`.")
     if name is not None:
         if labels is not None and "name" in labels:
             raise ValueError("`name` and the `name` key in `labels` cannot both be provided.")
@@ -617,6 +621,12 @@ def _create_job_spec(
     # make the job container reachable over SSH
     if ssh:
         job_spec["ssh"] = {"enabled": True}
+    # join a network group, optionally claiming aliases in it
+    if network_group:
+        network: dict[str, Any] = {"group": network_group}
+        if network_aliases:
+            network["aliases"] = network_aliases
+        job_spec["network"] = network
     # resource group is optional
     if resource_group_id:
         job_spec["resourceGroupId"] = resource_group_id

src/huggingface_hub/cli/jobs.py

@@ -216,6 +216,22 @@ def _parse_and_sync_job_volumes(
     ),
 ]
 
+NetworkGroupOpt = Annotated[
+    str | None,
+    Option(
+        "--network-group",
+        help="Join a network group. Jobs of the same owner sharing a group are placed together and reach each other on every port. Inside each member, `$HF_NETWORK_GROUP_HOSTNAME` resolves to every member. Lowercase alphanumerics and dashes, 46 characters max.",
+    ),
+]
+
+NetworkAliasOpt = Annotated[
+    list[str] | None,
+    Option(
+        "--network-alias",
+        help="Claim an alias in the network group. Members reach the jobs claiming it at `${HF_NETWORK_GROUP_PREFIX}<alias>`. Repeat the flag for several aliases. Requires `--network-group`.",
+    ),
+]
+
 WithOpt = Annotated[
     list[str] | None,
     Option(
@@ -353,6 +369,8 @@ def jobs_run(
     detach: DetachOpt = False,
     expose: ExposeOpt = None,
     ssh: SshEnabledOpt = False,
+    network_group: NetworkGroupOpt = None,
+    network_alias: NetworkAliasOpt = None,
     resource_group_id: ResourceGroupIdOpt = None,
     namespace: NamespaceOpt = None,
     token: TokenOpt = None,
@@ -373,6 +391,8 @@ def jobs_run(
         timeout=timeout,
         expose=expose,
         ssh=ssh,
+        network_group=network_group,
+        network_aliases=network_alias,
         resource_group_id=resource_group_id,
         namespace=namespace,
     )
@@ -388,6 +408,11 @@ def jobs_run(
         out.hint(f"Exposed ports are reachable at (requires an HF token with read access to the job):\n{urls}")
     if isinstance(job.status.ssh_url, str):
         out.hint(f"Use `hf jobs ssh {job.owner.name}/{job.id}` to open an SSH session into the job.")
+    if network_group:
+        out.hint(
+            f"Joined network group '{network_group}'. Jobs started with `--network-group {network_group}` reach each other "
+            "at `$HF_NETWORK_GROUP_HOSTNAME` (every member) or `${HF_NETWORK_GROUP_PREFIX}<alias>` (members claiming an alias)."
+        )
     if detach:
         job_ref = f"{job.owner.name}/{job.id}"
         out.hint(f"Use `hf jobs logs -f {job_ref}` to stream logs, or `hf jobs inspect {job_ref}` to check status.")
@@ -938,6 +963,8 @@ def jobs_uv_run(
     detach: DetachOpt = False,
     expose: ExposeOpt = None,
     ssh: SshEnabledOpt = False,
+    network_group: NetworkGroupOpt = None,
+    network_alias: NetworkAliasOpt = None,
     resource_group_id: ResourceGroupIdOpt = None,
     namespace: NamespaceOpt = None,
     token: TokenOpt = None,
@@ -963,6 +990,8 @@ def jobs_uv_run(
         timeout=timeout,
         expose=expose,
         ssh=ssh,
+        network_group=network_group,
+        network_aliases=network_alias,
         resource_group_id=resource_group_id,
         namespace=namespace,
     )
@@ -978,6 +1007,11 @@ def jobs_uv_run(
         out.hint(f"Exposed ports are reachable at (requires an HF token with read access to the job):\n{urls}")
     if isinstance(job.status.ssh_url, str):
         out.hint(f"Use `hf jobs ssh {job.owner.name}/{job.id}` to open an SSH session into the job.")
+    if network_group:
+        out.hint(
+            f"Joined network group '{network_group}'. Jobs started with `--network-group {network_group}` reach each other "
+            "at `$HF_NETWORK_GROUP_HOSTNAME` (every member) or `${HF_NETWORK_GROUP_PREFIX}<alias>` (members claiming an alias)."
+        )
     if detach:
         job_ref = f"{job.owner.name}/{job.id}"
         out.hint(f"Use `hf jobs logs -f {job_ref}` to stream logs, or `hf jobs inspect {job_ref}` to check status.")

src/huggingface_hub/hf_api.py

@@ -12130,6 +12130,8 @@ def run_job(
         volumes: list[Volume] | None = None,
         expose: list[int] | None = None,
         ssh: bool = False,
+        network_group: str | None = None,
+        network_aliases: list[str] | None = None,
         resource_group_id: str | None = None,
         namespace: str | None = None,
         token: bool | str | None = None,
@@ -12183,6 +12185,15 @@ def run_job(
                 write access to the job's namespace and an SSH public key registered on the Hub
                 (https://huggingface.co/settings/keys). Defaults to False.
 
+            network_group (`str`, *optional*):
+                Name of a network group to join. Jobs of the same owner sharing a group are placed together and
+                can reach each other on every port. Inside each member, `HF_NETWORK_GROUP_HOSTNAME` resolves to
+                every member of the group. Lowercase alphanumerics and dashes, 46 characters max.
+
+            network_aliases (`list[str]`, *optional*):
+                Aliases this job claims in its network group. Members reach the jobs claiming an alias at
+                `${HF_NETWORK_GROUP_PREFIX}<alias>`. Several jobs may claim the same alias. Requires `network_group`.
+
             resource_group_id (`str`, *optional*):
                 The ID of the resource group to create the Job in. Used to control access to resources within an
                 organization and for cost attribution/spending-limit features. If not provided, the Job is created
@@ -12241,6 +12252,8 @@ def run_job(
             volumes=volumes,
             expose=expose,
             ssh=ssh,
+            network_group=network_group,
+            network_aliases=network_aliases,
             resource_group_id=resource_group_id,
         )
         response = get_session().post(
@@ -12771,6 +12784,8 @@ def run_uv_job(
         volumes: list[Volume] | None = None,
         expose: list[int] | None = None,
         ssh: bool = False,
+        network_group: str | None = None,
+        network_aliases: list[str] | None = None,
         resource_group_id: str | None = None,
         namespace: str | None = None,
         token: bool | str | None = None,
@@ -12831,6 +12846,15 @@ def run_uv_job(
                 write access to the job's namespace and an SSH public key registered on the Hub
                 (https://huggingface.co/settings/keys). Defaults to False.
 
+            network_group (`str`, *optional*):
+                Name of a network group to join. Jobs of the same owner sharing a group are placed together and
+                can reach each other on every port. Inside each member, `HF_NETWORK_GROUP_HOSTNAME` resolves to
+                every member of the group. Lowercase alphanumerics and dashes, 46 characters max.
+
+            network_aliases (`list[str]`, *optional*):
+                Aliases this job claims in its network group. Members reach the jobs claiming an alias at
+                `${HF_NETWORK_GROUP_PREFIX}<alias>`. Several jobs may claim the same alias. Requires `network_group`.
+
             resource_group_id (`str`, *optional*):
                 The ID of the resource group to create the Job in. Used to control access to resources within an
                 organization and for cost attribution/spending-limit features. If not provided, the Job is created
@@ -12917,6 +12941,8 @@ def run_uv_job(
             volumes=volumes,
             expose=expose,
             ssh=ssh,
+            network_group=network_group,
+            network_aliases=network_aliases,
             resource_group_id=resource_group_id,
             namespace=namespace,
             token=token,

tests/test_cli.py

@@ -3353,6 +3353,8 @@ def test_run(self, runner: CliRunner) -> None:
             timeout=None,
             expose=None,
             ssh=False,
+            network_group=None,
+            network_aliases=None,
             resource_group_id=None,
             namespace=None,
         )
@@ -3381,6 +3383,8 @@ def test_run_with_extra_args(self, runner: CliRunner) -> None:
             timeout=None,
             expose=None,
             ssh=False,
+            network_group=None,
+            network_aliases=None,
             resource_group_id=None,
             namespace=None,
         )
@@ -3440,6 +3444,8 @@ def test_uv_command(self, runner: CliRunner) -> None:
             timeout=None,
             expose=None,
             ssh=False,
+            network_group=None,
+            network_aliases=None,
             resource_group_id=None,
             namespace=None,
         )
@@ -3471,6 +3477,8 @@ def test_uv_command_with_extra_args(self, runner: CliRunner) -> None:
             timeout=None,
             expose=None,
             ssh=False,
+            network_group=None,
+            network_aliases=None,
             resource_group_id=None,
             namespace=None,
         )
@@ -3500,6 +3508,8 @@ def test_uv_remote_script(self, runner: CliRunner) -> None:
             timeout=None,
             expose=None,
             ssh=False,
+            network_group=None,
+            network_aliases=None,
             resource_group_id=None,
             namespace=None,
         )
@@ -3530,6 +3540,8 @@ def test_uv_local_script(self, runner: CliRunner, tmp_path: Path) -> None:
             timeout=None,
             expose=None,
             ssh=False,
+            network_group=None,
+            network_aliases=None,
             resource_group_id=None,
             namespace=None,
         )
@@ -4451,6 +4463,42 @@ def test_serialize_expose(self, expose: list[int] | None, expected: dict | None)
         )
         assert spec.get("expose") == expected
 
+    @pytest.mark.parametrize(
+        "network_group, network_aliases, expected",
+        [
+            (None, None, None),
+            ("train", None, {"group": "train"}),
+            ("train", [], {"group": "train"}),
+            ("train", ["master", "worker"], {"group": "train", "aliases": ["master", "worker"]}),
+        ],
+    )
+    def test_serialize_network(
+        self, network_group: str | None, network_aliases: list[str] | None, expected: dict | None
+    ) -> None:
+        spec = _create_job_spec(
+            image="python:3.12",
+            command=["echo"],
+            env=None,
+            secrets=None,
+            flavor=None,
+            timeout=None,
+            network_group=network_group,
+            network_aliases=network_aliases,
+        )
+        assert spec.get("network") == expected
+
+    def test_network_aliases_require_group(self) -> None:
+        with pytest.raises(ValueError, match="network_aliases"):
+            _create_job_spec(
+                image="python:3.12",
+                command=["echo"],
+                env=None,
+                secrets=None,
+                flavor=None,
+                timeout=None,
+                network_aliases=["master"],
+            )
+
 
 class TestWebhooksCommand:
     def _make_webhook(self, **kwargs):