Skip to content

Commit

Permalink
Remove try/finally in pytest fixture to simplify PR diff
Browse files Browse the repository at this point in the history
  • Loading branch information
@ml-evs
ml-evs committed Nov 25, 2024
1 parent 4476ce3 commit 02dbef9
Showing 1 changed file with 154 additions and 147 deletions.
301 changes: 154 additions & 147 deletions tests/integration/conftest.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -225,161 +225,168 @@ def write_tmp_settings(
original_jf_remote_projects_folder = os.environ.get("JFREMOTE_PROJECTS_FOLDER")
original_jf_remote_project = os.environ.get("JFREMOTE_PROJECT")
original_config_file = os.environ.get("JFREMOTE_CONFIG_FILE")
try:
os.environ["JFREMOTE_PROJECTS_FOLDER"] = str(tmp_dir.resolve())
workdir = tmp_dir / "jfr"
workdir.mkdir(exist_ok=True)
os.environ["JFREMOTE_PROJECT"] = random_project_name
# Set config file to a random path so that we don't accidentally load the default
os.environ["JFREMOTE_CONFIG_FILE"] = _get_random_name(length=10) + ".json"
# This import must come after setting the env vars as jobflow loads the default
# config on import
from jobflow_remote.config import Project

project = Project(
name=random_project_name,
jobstore={
"docs_store": {
"type": "MongoStore",
"database": store_database_name,
"host": "localhost",
"port": db_port,
"collection_name": "docs",
},
"additional_stores": {
"big_data": {
"type": "GridFSStore",
"database": store_database_name,
"host": "localhost",
"port": db_port,
"collection_name": "data",
},
},

os.environ["JFREMOTE_PROJECTS_FOLDER"] = str(tmp_dir.resolve())
workdir = tmp_dir / "jfr"
workdir.mkdir(exist_ok=True)
os.environ["JFREMOTE_PROJECT"] = random_project_name
# Set config file to a random path so that we don't accidentally load the default
os.environ["JFREMOTE_CONFIG_FILE"] = _get_random_name(length=10) + ".json"
# This import must come after setting the env vars as jobflow loads the default
# config on import
from jobflow_remote.config import Project

project = Project(
name=random_project_name,
jobstore={
"docs_store": {
"type": "MongoStore",
"database": store_database_name,
"host": "localhost",
"port": db_port,
"collection_name": "docs",
},
queue={
"store": {
"type": "MongoStore",
"additional_stores": {
"big_data": {
"type": "GridFSStore",
"database": store_database_name,
"host": "localhost",
"port": db_port,
"collection_name": "jobs",
"collection_name": "data",
},
"flows_collection": "flows",
},
log_level="debug",
workers={
"test_local_worker": dict(
type="local",
scheduler_type="shell",
work_dir=str(workdir),
resources={},
),
"test_remote_slurm_worker": dict(
type="remote",
host="localhost",
port=slurm_ssh_port,
scheduler_type="slurm",
work_dir="/home/jobflow/jfr",
user="jobflow",
password="jobflow",
pre_run="source /home/jobflow/.venv/bin/activate",
resources={"partition": "debug", "ntasks": 1, "time": "00:01:00"},
connect_kwargs={"allow_agent": False, "look_for_keys": False},
),
"test_remote_sge_worker": dict(
type="remote",
host="localhost",
port=sge_ssh_port,
scheduler_type="sge",
work_dir="/home/jobflow/jfr",
user="jobflow",
password="jobflow",
scheduler_username="jobflow",
pre_run="source /home/jobflow/.venv/bin/activate",
resources={"partition": "debug", "ntasks": 1, "time": "00:01:00"},
connect_kwargs={"allow_agent": False, "look_for_keys": False},
),
"test_batch_remote_worker": dict(
type="remote",
host="localhost",
port=slurm_ssh_port,
scheduler_type="slurm",
work_dir="/home/jobflow/jfr",
user="jobflow",
password="jobflow",
pre_run="source /home/jobflow/.venv/bin/activate",
resources={"partition": "debug", "ntasks": 1, "time": "00:01:00"},
connect_kwargs={"allow_agent": False, "look_for_keys": False},
batch={
"jobs_handle_dir": "/home/jobflow/jfr/batch_handle",
"work_dir": "/home/jobflow/jfr/batch_work",
"max_wait": 10,
},
max_jobs=1,
),
"test_max_jobs_worker": dict(
type="local",
scheduler_type="shell",
work_dir=str(workdir),
resources={},
max_jobs=2,
),
"test_batch_multi_remote_worker": dict(
type="remote",
host="localhost",
port=slurm_ssh_port,
scheduler_type="slurm",
work_dir="/home/jobflow/jfr",
user="jobflow",
password="jobflow",
pre_run="source /home/jobflow/.venv/bin/activate",
resources={"partition": "debug", "ntasks": 1, "time": "00:01:00"},
connect_kwargs={"allow_agent": False, "look_for_keys": False},
batch={
"jobs_handle_dir": "/home/jobflow/jfr/batch_multi_handle",
"work_dir": "/home/jobflow/jfr/batch_multi_work",
"max_wait": 10,
"parallel_jobs": 2,
},
max_jobs=1,
),
"test_sanitize_remote_worker": dict(
type="remote",
host="localhost",
port=slurm_ssh_port,
scheduler_type="slurm",
work_dir="/home/jobflow/jfr",
user="jobflow",
password="jobflow",
pre_run="source /home/jobflow/.venv/bin/activate",
resources={"partition": "debug", "ntasks": 1, "time": "00:01:00"},
connect_kwargs={"allow_agent": False, "look_for_keys": False},
sanitize_command=True,
),
},
queue={
"store": {
"type": "MongoStore",
"database": store_database_name,
"host": "localhost",
"port": db_port,
"collection_name": "jobs",
},
exec_config={"test": {"export": {"TESTING_ENV_VAR": random_project_name}}},
runner=dict(
delay_checkout=1,
delay_check_run_status=1,
delay_advance_status=1,
max_step_attempts=3,
delta_retry=(1, 1, 1),
"flows_collection": "flows",
},
log_level="debug",
workers={
"test_sanitize_local_worker": dict(
type="local",
scheduler_type="shell",
work_dir=str(workdir),
resources={},
sanitize_command=True,
),
)
project_json = project.model_dump_json(indent=2)
with open(tmp_dir / f"{random_project_name}.json", "w") as f:
f.write(project_json)

yield project
finally:
shutil.rmtree(tmp_dir)
# Reset environment variables if they were set elsewhere
if original_jf_remote_projects_folder is not None:
os.environ["JFREMOTE_PROJECTS_FOLDER"] = original_jf_remote_projects_folder
if original_jf_remote_project is not None:
os.environ["JFREMOTE_PROJECT"] = original_jf_remote_project
if original_config_file is not None:
os.environ["JFREMOTE_CONFIG_FILE"] = original_config_file
"test_local_worker": dict(
type="local",
scheduler_type="shell",
work_dir=str(workdir),
resources={},
),
"test_remote_slurm_worker": dict(
type="remote",
host="localhost",
port=slurm_ssh_port,
scheduler_type="slurm",
work_dir="/home/jobflow/jfr",
user="jobflow",
password="jobflow",
pre_run="source /home/jobflow/.venv/bin/activate",
resources={"partition": "debug", "ntasks": 1, "time": "00:01:00"},
connect_kwargs={"allow_agent": False, "look_for_keys": False},
),
"test_remote_sge_worker": dict(
type="remote",
host="localhost",
port=sge_ssh_port,
scheduler_type="sge",
work_dir="/home/jobflow/jfr",
user="jobflow",
password="jobflow",
scheduler_username="jobflow",
pre_run="source /home/jobflow/.venv/bin/activate",
resources={"partition": "debug", "ntasks": 1, "time": "00:01:00"},
connect_kwargs={"allow_agent": False, "look_for_keys": False},
),
"test_batch_remote_worker": dict(
type="remote",
host="localhost",
port=slurm_ssh_port,
scheduler_type="slurm",
work_dir="/home/jobflow/jfr",
user="jobflow",
password="jobflow",
pre_run="source /home/jobflow/.venv/bin/activate",
resources={"partition": "debug", "ntasks": 1, "time": "00:01:00"},
connect_kwargs={"allow_agent": False, "look_for_keys": False},
batch={
"jobs_handle_dir": "/home/jobflow/jfr/batch_handle",
"work_dir": "/home/jobflow/jfr/batch_work",
"max_wait": 10,
},
max_jobs=1,
),
"test_batch_multi_remote_worker": dict(
type="remote",
host="localhost",
port=slurm_ssh_port,
scheduler_type="slurm",
work_dir="/home/jobflow/jfr",
user="jobflow",
password="jobflow",
pre_run="source /home/jobflow/.venv/bin/activate",
resources={"partition": "debug", "ntasks": 1, "time": "00:01:00"},
connect_kwargs={"allow_agent": False, "look_for_keys": False},
batch={
"jobs_handle_dir": "/home/jobflow/jfr/batch_multi_handle",
"work_dir": "/home/jobflow/jfr/batch_multi_work",
"max_wait": 10,
"parallel_jobs": 2,
},
max_jobs=1,
),
"test_max_jobs_worker": dict(
type="local",
scheduler_type="shell",
work_dir=str(workdir),
resources={},
max_jobs=2,
),
"test_sanitize_remote_worker": dict(
type="remote",
host="localhost",
port=slurm_ssh_port,
scheduler_type="slurm",
work_dir="/home/jobflow/jfr",
user="jobflow",
password="jobflow",
pre_run="source /home/jobflow/.venv/bin/activate",
resources={"partition": "debug", "ntasks": 1, "time": "00:01:00"},
connect_kwargs={"allow_agent": False, "look_for_keys": False},
sanitize_command=True,
),
},
exec_config={"test": {"export": {"TESTING_ENV_VAR": random_project_name}}},
runner=dict(
delay_checkout=1,
delay_check_run_status=1,
delay_advance_status=1,
max_step_attempts=3,
delta_retry=(1, 1, 1),
),
)
project_json = project.model_dump_json(indent=2)
with open(tmp_dir / f"{random_project_name}.json", "w") as f:
f.write(project_json)

yield project

shutil.rmtree(tmp_dir)
# Reset environment variables if they were set elsewhere
if original_jf_remote_projects_folder is not None:
os.environ["JFREMOTE_PROJECTS_FOLDER"] = original_jf_remote_projects_folder
if original_jf_remote_project is not None:
os.environ["JFREMOTE_PROJECT"] = original_jf_remote_project
if original_config_file is not None:
os.environ["JFREMOTE_CONFIG_FILE"] = original_config_file


@pytest.fixture()
Expand Down

0 comments on commit 02dbef9

Please sign in to comment.