add: diffusers package to the hf stack

change: pinned and upgraded packages from hf stack
fix: telemetry test

Signed-off-by: DWarez <dario.salvati@huggingface.co>

Author: dwarez

docker/huggingface/pytorch/Dockerfile

@@ -12,10 +12,17 @@ ARG DLC_MAJOR_VERSION=1
 LABEL maintainer="Hugging Face"
 LABEL dlc_major_version="${DLC_MAJOR_VERSION}"
 
-ARG TRANSFORMERS_VERSION=5.10.2
-ARG HUGGINGFACE_HUB_VERSION=1.17.0
-ARG HF_XET_VERSION=1.5.0
+ARG TRANSFORMERS_VERSION=5.13.1
+ARG HUGGINGFACE_HUB_VERSION=1.23.0
+ARG HF_XET_VERSION=1.5.1
 ARG TRL_VERSION=1.8.0
+ARG DIFFUSERS_VERSION=0.39.0
+ARG HF_TRANSFER_VERSION=0.1.9
+ARG DATASETS_VERSION=5.0.0
+ARG ACCELERATE_VERSION=1.14.0
+ARG PEFT_VERSION=0.19.1
+ARG EVALUATE_VERSION=0.4.6
+ARG MATH_VERIFY_VERSION=0.9.0
 # vllm >= 0.24 required: earlier versions ship CVE-affected deps (mooncake
 # etc.). trl's [vllm] extra caps vllm at 0.23, so trl is installed without it —
 # vllm itself provides the extra's remaining deps (fastapi, pydantic, uvicorn).
@@ -36,11 +43,15 @@ RUN uv pip install --python ${VENV_PYTHON} \
   transformers==${TRANSFORMERS_VERSION} \
   huggingface-hub==${HUGGINGFACE_HUB_VERSION} \
   hf-xet==${HF_XET_VERSION} \
-  hf-transfer \
-  datasets accelerate peft evaluate \
+  "hf-transfer==${HF_TRANSFER_VERSION}" \
+  "datasets==${DATASETS_VERSION}" \
+  "accelerate==${ACCELERATE_VERSION}" \
+  "peft==${PEFT_VERSION}" \
+  "evaluate==${EVALUATE_VERSION}" \
+  "diffusers==${DIFFUSERS_VERSION}" \
   "trl==${TRL_VERSION}" \
   "vllm==${VLLM_VERSION}" \
-  math-verify \
+  "math-verify==${MATH_VERIFY_VERSION}" \
   sentencepiece protobuf einops scipy scikit-learn tensorboard \
   "soupsieve>=2.8.4" \
   && uv cache clean

test/telemetry/test_telemetry.py

@@ -24,6 +24,7 @@
 """
 
 import logging
+import time
 from urllib.parse import parse_qs, urlparse
 
 import pytest
@@ -40,6 +41,28 @@ def delete_telemetry_tag(aws_session, instance_id):
     aws_session.ec2.delete_tags(Resources=[instance_id], Tags=[{"Key": TELEMETRY_TAG_KEY}])
 
 
+def set_imds_hop_limit(aws_session, instance_id, hop_limit, timeout=120):
+    """Set the IMDSv2 hop limit and wait until it is live.
+
+    modify_instance_metadata_options is asynchronous (InstanceMetadataOptions.State
+    goes pending -> applied); launching a container before it applies races the
+    telemetry script's one-shot IMDS probe and makes the test flaky.
+    """
+    aws_session.ec2.modify_instance_metadata_options(
+        InstanceId=instance_id, HttpPutResponseHopLimit=hop_limit
+    )
+    deadline = time.time() + timeout
+    while time.time() < deadline:
+        instance = aws_session.ec2.describe_instances(InstanceIds=[instance_id])["Reservations"][0][
+            "Instances"
+        ][0]
+        options = instance["InstanceMetadataOptions"]
+        if options["State"] == "applied" and options["HttpPutResponseHopLimit"] == hop_limit:
+            return
+        time.sleep(5)
+    raise TimeoutError(f"IMDS hop_limit={hop_limit} not applied within {timeout}s")
+
+
 def telemetry_file_exists(conn, container):
     result = conn.run(f"docker exec {container} test -f /tmp/test_request.txt", warn=True)
     return result.return_code == 0
@@ -122,9 +145,7 @@ def test_imds_hop_limit(conn, aws_session, telemetry_container, ec2_instance):
 
     # hop_limit=1: container can't get IMDSv2 token through Docker bridge
     delete_telemetry_tag(aws_session, instance_id)
-    aws_session.ec2.modify_instance_metadata_options(
-        InstanceId=instance_id, HttpPutResponseHopLimit=1
-    )
+    set_imds_hop_limit(aws_session, instance_id, 1)
     container_name = telemetry_container("telemetry-hop1", env={"TEST_MODE": "1"})
 
     assert not telemetry_file_exists(conn, container_name), (
@@ -135,9 +156,7 @@ def test_imds_hop_limit(conn, aws_session, telemetry_container, ec2_instance):
     assert TELEMETRY_TAG_KEY not in tags, "Tag should not exist with hop_limit=1"
 
     # hop_limit=2: container can get IMDSv2 token
-    aws_session.ec2.modify_instance_metadata_options(
-        InstanceId=instance_id, HttpPutResponseHopLimit=2
-    )
+    set_imds_hop_limit(aws_session, instance_id, 2)
     container_name = telemetry_container("telemetry-hop2", env={"TEST_MODE": "1"})
 
     assert telemetry_file_exists(conn, container_name), "Telemetry should fire with hop_limit=2"