Skip to content

Commit 59e9145

Browse files
committed
Change sdk to run
Signed-off-by: Ryan Wolf <[email protected]>
1 parent 1050604 commit 59e9145

File tree

3 files changed

+9
-9
lines changed

3 files changed

+9
-9
lines changed

examples/nemo_sdk/launch_slurm.py renamed to examples/nemo_run/launch_slurm.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,13 +12,13 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
import nemo_sdk as sdk
16-
from nemo_sdk.core.execution import SlurmExecutor
15+
import nemo_run as run
16+
from nemo_run.core.execution import SlurmExecutor
1717

18-
from nemo_curator.nemo_sdk import SlurmJobConfig
18+
from nemo_curator.nemo_run import SlurmJobConfig
1919

2020

21-
@sdk.factory
21+
@run.factory
2222
def nemo_curator_slurm_executor() -> SlurmExecutor:
2323
"""
2424
Configure the following function with the details of your SLURM cluster
File renamed without changes.

nemo_curator/nemo_sdk/slurm.py renamed to nemo_curator/nemo_run/slurm.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,14 @@
1717

1818
from 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
2424
class 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

Comments
 (0)