Remove from mcli.sdk imports (#793)

mosaicml · Dec 11, 2023 · 693ef10 · 693ef10
1 parent 8d96f9d
commit 693ef10
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 10 deletions.
diff --git a/.github/mcp/mcp_pytest.py b/.github/mcp/mcp_pytest.py
@@ -6,8 +6,8 @@
 import argparse
 import time
 
-from mcli.sdk import (RunConfig, RunStatus, create_run, follow_run_logs,
-                      wait_for_run_status)
+from mcli import (RunConfig, RunStatus, create_run, follow_run_logs,
+                  wait_for_run_status)
 
 if __name__ == '__main__':
 
@@ -107,9 +107,11 @@
 
     config = RunConfig(
         name=name,
-        cluster=args.cluster,
-        gpu_type=args.gpu_type,
-        gpu_num=args.gpu_num,
+        compute={
+            'cluster': args.cluster,
+            'gpu_type': args.gpu_type,
+            'gpus': args.gpu_num
+        },
         image=args.image,
         integrations=[git_integration],
         command=command,

diff --git a/scripts/train/benchmarking/submit_benchmarks.py b/scripts/train/benchmarking/submit_benchmarks.py
@@ -7,8 +7,8 @@
 
 import requests
 import yaml
-from mcli.models.run_config import SchedulingConfig
-from mcli.sdk import RunConfig, create_run, get_clusters
+
+from mcli import RunConfig, SchedulingConfig, create_run, get_clusters
 
 
 def _get_cluster_info():
@@ -470,9 +470,11 @@ def run_config(config: Tuple[str, int, int, str, str, int, str],
         parameters['model']['fc_type'] = 'te'
     # Create run config mcli sdk/api
     config = RunConfig(name=name,
-                       gpu_type=gpu_type,
-                       gpu_num=gpu_num,
-                       cluster=cluster,
+                       compute={
+                           'cluster': cluster,
+                           'gpu_type': gpu_type,
+                           'gpus': gpu_num
+                       },
                        image=args.image,
                        integrations=integrations,
                        command=command,