1. Generating GCG Suffixes Using Azure Machine Learning

import os

from pyrit.setup.initialization import _load_environment_files

_load_environment_files(env_files=None)

subscription_id = os.environ.get("AZURE_ML_SUBSCRIPTION_ID")
resource_group = os.environ.get("AZURE_ML_RESOURCE_GROUP")
workspace = os.environ.get("AZURE_ML_WORKSPACE_NAME")
print(workspace)

import logging
import warnings

logging.getLogger("azure.ai.ml").setLevel(logging.ERROR)
warnings.filterwarnings("ignore", module=r"azure\.ai\.ml.*")

from azure.ai.ml import MLClient
from azure.identity import AzureCliCredential

ml_client = MLClient(AzureCliCredential(), subscription_id, resource_group, workspace)

from pathlib import Path

from azure.ai.ml.entities import BuildContext, Environment

from pyrit.common.path import HOME_PATH

# Configure the AML environment — build context is the repo root so the Dockerfile
# can COPY pyproject.toml and pyrit/ for pip install -e ".[gcg]"
env_docker_context = Environment(
    build=BuildContext(
        path=Path(HOME_PATH),
        dockerfile_path="pyrit/auxiliary_attacks/gcg/src/Dockerfile",
    ),
    name="pyrit-gcg",
    description="PyRIT GCG environment: CUDA 12.1 + Python 3.11 + pip install -e .[gcg]",
    tags={"Owner": os.environ.get("USER", "unknown")},
)

ml_client.environments.create_or_update(env_docker_context)

import tempfile

from pyrit.auxiliary_attacks.gcg import (
    GCGAlgorithmConfig,
    GCGConfig,
    GCGDataConfig,
    GCGModelConfig,
    GCGOutputConfig,
)

config = GCGConfig(
    models=[GCGModelConfig(name="meta-llama/Llama-2-7b-chat-hf")],
    algorithm=GCGAlgorithmConfig(n_steps=5, batch_size=64, test_steps=1),
    output=GCGOutputConfig(result_prefix="gcg_suffix"),
)
data_config = GCGDataConfig(
    train_data=("https://raw.githubusercontent.com/llm-attacks/llm-attacks/main/data/advbench/harmful_behaviors.csv"),
    n_train_data=5,
    n_test_data=0,
)

# Write the configs into a tempdir so AML can mount them as separate job inputs.
config_dir = Path(tempfile.mkdtemp(prefix="gcg-aml-config-"))
config_path = config_dir / "config.json"
data_path = config_dir / "data.json"
config.to_json_file(config_path)
data_config.to_json_file(data_path)

from azure.ai.ml import Input, Output, command

job = command(
    code=Path(HOME_PATH),
    command=(
        "python -m pyrit.auxiliary_attacks.gcg.experiments.run"
        " --config ${{inputs.config}}"
        " --data ${{inputs.data}}"
        " --output-dir ${{outputs.results}}"
    ),
    inputs={
        "config": Input(type="uri_file", path=str(config_path)),
        "data": Input(type="uri_file", path=str(data_path)),
    },
    outputs={"results": Output(type="uri_folder")},
    environment=f"{env_docker_context.name}:{env_docker_context.version}",
    environment_variables={"HUGGINGFACE_TOKEN": os.environ["HUGGINGFACE_TOKEN"]},
    compute="gcg-gpu-a100",
    display_name="gcg_suffix_generation",
    description="Generate adversarial suffixes using GCG on Llama-2.",
    tags={"Owner": os.environ.get("USER", "unknown")},
)

returned_job = ml_client.create_or_update(job)
print(f"Job: {returned_job.name}")
print(f"Status: {returned_job.status}")
print(f"Studio URL: {returned_job.studio_url}")

import json
import tempfile
import time
from pathlib import Path

_TERMINAL_STATES = {"Completed", "Failed", "Canceled", "CancelRequested"}

last_status = None
while True:
    current_status = ml_client.jobs.get(returned_job.name).status
    if current_status != last_status:
        print(f"Job status: {current_status}", flush=True)
        last_status = current_status
    if current_status in _TERMINAL_STATES:
        break
    time.sleep(60)

assert current_status == "Completed", f"Job did not complete successfully: {current_status}"

download_dir = Path(tempfile.mkdtemp(prefix="gcg-aml-"))
ml_client.jobs.download(name=returned_job.name, download_path=str(download_dir), all=True)

result_files = list(download_dir.rglob("gcg_suffix_*.json"))
if not result_files:
    print(f"No GCG result file found under {download_dir}. Files captured:")
    for p in sorted(download_dir.rglob("*")):
        if p.is_file():
            print(f"  {p.relative_to(download_dir)}")
    raise FileNotFoundError("Result JSON not in downloaded artifacts")

result_file = result_files[0]
with open(result_file) as f:
    log = json.load(f)

final_suffix = log["controls"][-1] if log["controls"] else None
final_loss = log["losses"][-1] if log["losses"] else None

print(f"Result file: {result_file.name}")
print(f"Steps run: {len(log['controls'])}")
print(f"Final loss: {final_loss}")
print(f"Generated suffix: {final_suffix!r}")