1717
1818from nemo_curator .utils .import_utils import safe_import
1919
20- sdk = safe_import ("nemo_sdk " )
20+ run = safe_import ("nemo_run " )
2121
2222
2323@dataclass
2424class SlurmJobConfig :
2525 """
2626 Configuration for running a NeMo Curator script on a SLURM cluster using
27- NeMo SDK
27+ NeMo Run
2828
2929 Args:
3030 job_dir: The base directory where all the files related to setting up
@@ -69,7 +69,7 @@ class SlurmJobConfig:
6969
7070 def to_script (self , add_scheduler_file : bool = True , add_device : bool = True ):
7171 """
72- Converts to a script object executable by NeMo SDK
72+ Converts to a script object executable by NeMo Run
7373 Args:
7474 add_scheduler_file: Automatically appends a '--scheduler-file' argument to the
7575 script_command where the value is job_dir/logs/scheduler.json. All
@@ -79,7 +79,7 @@ def to_script(self, add_scheduler_file: bool = True, add_device: bool = True):
7979 where the value is the member variable of device. All scripts included in
8080 NeMo Curator accept and require this argument.
8181 Returns:
82- A NeMo SDK Script that will intialize a Dask cluster, and run the specified command.
82+ A NeMo Run Script that will intialize a Dask cluster, and run the specified command.
8383 It is designed to be executed on a SLURM cluster
8484 """
8585 env_vars = self ._build_env_vars ()
@@ -94,7 +94,7 @@ def to_script(self, add_scheduler_file: bool = True, add_device: bool = True):
9494 # Surround the command in quotes so the variable gets set properly
9595 env_vars ["SCRIPT_COMMAND" ] = f"\" { env_vars ['SCRIPT_COMMAND' ]} \" "
9696
97- return sdk .Script (path = self .container_entrypoint , env = env_vars )
97+ return run .Script (path = self .container_entrypoint , env = env_vars )
9898
9999 def _build_env_vars (self ) -> Dict [str , str ]:
100100 env_vars = vars (self )
0 commit comments