def download_model_from_task()

in utils/run_model.py [0:0]

• 45 lines of code
• 14 McCabe index (conditional complexity)

def download_model_from_task(output: str, task_id: str) -> str:
    """Downloads to the output directory in a subfolder {src}-{trg}-{task_id}"""
    options = {"rootUrl": ("%s" % TC_MOZILLA)}
    queue: Any = taskcluster.Queue(options=options)
    task = queue.task(task_id)
    status = queue.status(task_id)

    if status["status"]["state"] != "completed":
        raise Exception("The task was not completed")

    task_name = task["metadata"]["name"]
    language_pair = get_language_pair_from_task_name(task_name)
    if not language_pair:
        raise Exception(f'Could not find the language pair for the task "{task_name}"')

    print(f"Downloading from task {task_name}")

    artifacts = queue.listLatestArtifacts(task_id)["artifacts"]

    for artifact in artifacts:
        print(artifact["name"])

    decoder = next(
        a for a in artifacts if a["name"].endswith("/final.model.npz.best-chrf.npz.decoder.yml")
    )
    model = next(a for a in artifacts if a["name"].endswith("/final.model.npz.best-chrf.npz"))
    vocab = next(a for a in artifacts if a["name"].endswith("/vocab.spm"))

    if not decoder:
        raise Exception("Could not find the decoder in artifacts")
    if not model:
        raise Exception("Could not find the model in artifacts")
    if not vocab:
        raise Exception("Could not find the vocab in artifacts")

    model_path = os.path.join(output, f"{language_pair}-{task_id}")
    print(model_path)

    os.makedirs(model_path, exist_ok=True)

    print(
        f'Downloading models from "{task_name}": '
        f"https://firefox-ci-tc.services.mozilla.com/tasks/{task_id}"
    )

    downloads = [
        (decoder, "decoder.yml"),
        (model, "model.npz"),
        (vocab, "vocab.spm"),
    ]

    for artifact, filename in downloads:
        stream_download_to_file(
            queue.buildUrl("getLatestArtifact", task_id, artifact["name"]),
            os.path.join(model_path, filename),
        )

    print(f"Model files are available at: {model_path}")
    return model_path