Skip to content

Commit

Permalink
Wrap train-taskcluster.sh in train_taskcluster.py
Browse files Browse the repository at this point in the history
  • Loading branch information
bhearsum committed May 7, 2024
1 parent 5c29750 commit 43bbc78
Show file tree
Hide file tree
Showing 8 changed files with 106 additions and 5 deletions.
2 changes: 1 addition & 1 deletion Taskfile.yml
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@ tasks:
task test -- tests/test_alignments.py
deps: [poetry-install-tests]
cmds:
- PYTHONPATH=$(pwd) poetry run pytest -vv {{.CLI_ARGS}}
- PYTHONPATH="$(pwd):$(pwd)/taskcluster/scripts/pipeline" poetry run pytest -vv {{.CLI_ARGS}}

test-docker:
desc: Run the unit tests in the docker image. Some tests require the pre-built Linux executables.
Expand Down
4 changes: 3 additions & 1 deletion taskcluster/kinds/finetune-student/kind.yml
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,8 @@ tasks:
- pipeline/train/configs/opustrainer/student.yml
- pipeline/train/configs/training/student.train.yml
- pipeline/train/train.sh
- taskcluster/scripts/pipeline/train_taskcluster.py
- taskcluster/scripts/pipeline/train-taskcluster.sh
from-parameters:
marian_args: training_config.marian-args.training-student-finetuned

Expand Down Expand Up @@ -93,7 +95,7 @@ tasks:
pip3 install $VCS_PATH/tracking &&
export PATH="$HOME/.local/bin:$PATH" &&
export MARIAN=$MOZ_FETCHES_DIR &&
$VCS_PATH/taskcluster/scripts/pipeline/train-taskcluster.sh
$VCS_PATH/taskcluster/scripts/pipeline/train_taskcluster.py
student
finetune
{src_locale}
Expand Down
3 changes: 2 additions & 1 deletion taskcluster/kinds/train-backwards/kind.yml
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ tasks:
type: train-backwards
resources:
- pipeline/train/train.sh
- taskcluster/scripts/pipeline/train_taskcluster.py
- taskcluster/scripts/pipeline/train-taskcluster.sh
- pipeline/train/configs/model/backward.yml
- pipeline/train/configs/opustrainer/backward.yml
Expand Down Expand Up @@ -98,7 +99,7 @@ tasks:
pip3 install $VCS_PATH/tracking &&
export PATH="$HOME/.local/bin:$PATH" &&
export MARIAN=$MOZ_FETCHES_DIR &&
$VCS_PATH/taskcluster/scripts/pipeline/train-taskcluster.sh
$VCS_PATH/taskcluster/scripts/pipeline/train_taskcluster.py
backward
train
{trg_locale}
Expand Down
4 changes: 3 additions & 1 deletion taskcluster/kinds/train-student/kind.yml
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,8 @@ tasks:
- pipeline/train/configs/opustrainer/student.yml
- pipeline/train/configs/training/student.train.yml
- pipeline/train/train.sh
- taskcluster/scripts/pipeline/train_taskcluster.py
- taskcluster/scripts/pipeline/train-taskcluster.sh
from-parameters:
marian_args: training_config.marian-args.training-student
worker-type:
Expand Down Expand Up @@ -94,7 +96,7 @@ tasks:
pip3 install $VCS_PATH/tracking &&
export PATH="$HOME/.local/bin:$PATH" &&
export MARIAN=$MOZ_FETCHES_DIR &&
$VCS_PATH/taskcluster/scripts/pipeline/train-taskcluster.sh
$VCS_PATH/taskcluster/scripts/pipeline/train_taskcluster.py
student
train
{src_locale}
Expand Down
3 changes: 2 additions & 1 deletion taskcluster/kinds/train-teacher/kind.yml
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ tasks:
- pipeline/train/configs/opustrainer/teacher.yml
- pipeline/train/configs/training/teacher.train.yml
- pipeline/train/train.sh
- taskcluster/scripts/pipeline/train_taskcluster.py
- taskcluster/scripts/pipeline/train-taskcluster.sh
from-parameters:
marian_args: training_config.marian-args.training-teacher
Expand Down Expand Up @@ -119,7 +120,7 @@ tasks:
pip3 install $VCS_PATH/tracking &&
export PATH="$HOME/.local/bin:$PATH" &&
export MARIAN=$MOZ_FETCHES_DIR &&
$VCS_PATH/taskcluster/scripts/pipeline/train-taskcluster.sh
$VCS_PATH/taskcluster/scripts/pipeline/train_taskcluster.py
teacher
train
{src_locale}
Expand Down
Empty file.
15 changes: 15 additions & 0 deletions taskcluster/scripts/pipeline/train_taskcluster.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
#!/usr/bin/env python3

import os.path
import subprocess
import sys

TRAINING_SCRIPT = os.path.join(os.path.dirname(__file__), "train-taskcluster.sh")


def main(args):
subprocess.run([TRAINING_SCRIPT, *args], check=True)


if __name__ == "__main__":
main(sys.argv[1:])
80 changes: 80 additions & 0 deletions tests/test_train_taskcluster.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
import os
from unittest import mock

import pytest
import train_taskcluster

TRAIN_TASKCLUSTER_SH = os.path.normpath(
os.path.join(
os.path.dirname(os.path.realpath(__file__)),
"..",
"taskcluster",
"scripts",
"pipeline",
"train-taskcluster.sh",
)
)


@pytest.mark.parametrize(
"args",
(
pytest.param(
[
"model_type",
"type",
"src",
"trg",
"train_set_prefix",
"valid_set_prefix",
"model_dir",
"best_model_metric",
"alignments",
"seed",
],
id="required_only",
),
pytest.param(
[
"model_type",
"type",
"src",
"trg",
"train_set_prefix",
"valid_set_prefix",
"model_dir",
"best_model_metric",
"alignments",
"seed",
"pretrained_model_mode",
"pretrained_model_type",
],
id="with_pretrained_model",
),
pytest.param(
[
"model_type",
"type",
"src",
"trg",
"train_set_prefix",
"valid_set_prefix",
"model_dir",
"best_model_metric",
"alignments",
"seed",
"pretrained_model_mode",
"pretrained_model_type",
"--foo",
"--bar",
],
id="with_extra_params",
),
),
)
def test_all_args_forwarded(args):
with mock.patch("train_taskcluster.subprocess") as mocked_subprocess:
train_taskcluster.main(args)
assert mocked_subprocess.run.call_args_list == [
mock.call([TRAIN_TASKCLUSTER_SH] + args, check=True),
]

0 comments on commit 43bbc78

Please sign in to comment.