manual deploy with tag + nbox fail fast

Author: yashbonde

nbox/assets/exe.jinja

@@ -10,7 +10,6 @@ os.environ["PYTHONUNBUFFERED"] = "true" # so print comes when it should come
 import fire
 import inspect
 from functools import lru_cache, partial
-from {{ file_name }} import {{ fn_name }}
 
 from nbox import Operator, logger
 import nbox.utils as U
@@ -20,6 +19,11 @@ from nbox.lib.dist import NBXLet
 def get_op(cloud = False) -> Operator:
   # The beauty of this function is that it ensures that the operator class is loaded only once
   try:
+    # import user code, add this to the try/except because if the code does not exit and there
+    # is an infinite loop, there can be a whole bunch of side effects, ex: 100s of LMAO live trackers
+    from {{ file_name }} import {{ fn_name }}
+    
+    # load the operator
     obj = {{ init_code }}
     if not type(obj) == Operator and {{ load_operator }}:
       # there is an initial level of precaution that we use during deployment, but we are adding simple

nbox/auth.py

@@ -213,6 +213,9 @@ def init_secret():
       token_present = len(secret.access_token) > 0,
       nbx_url = secret.nbx_url,
     ))
+
+    if not secret.workspace_id:
+      raise Exception("Workspace ID not found. Please run `nbox login` to login to NimbleBox.")
     return secret
   else:
     logger.info(f"Skipping authentication as NBOX_NO_AUTH is set to True")

nbox/init.py

@@ -22,7 +22,7 @@
 
 
 from nbox.auth import secret, AuthConfig
-from nbox.utils import logger, env
+from nbox.utils import logger, env, hard_exit_program
 from nbox.subway import Sub30
 from nbox.hyperloop.jobs.nbox_ws_pb2_grpc import WSJobServiceStub
 from nbox.hyperloop.deploy.serve_pb2_grpc import ServingServiceStub, ModelServiceStub
@@ -110,6 +110,7 @@ def create_webserver_subway(version: str = "v1", session: requests.Session = Non
   except Exception as e:
     logger.error(f"Could not connect to webserver at {secret('nbx_url')}")
     logger.error(e)
+    
     return None
 
   spec = r.json()

nbox/jobs.py

@@ -17,7 +17,8 @@
 from nbox.auth import secret, AuthConfig, auth_info_pb
 from nbox.utils import logger
 from nbox.version import __version__
-from nbox.messages import rpc, streaming_rpc
+from nbox import messages as mpb
+# from nbox.messages import rpc, streaming_rpc
 from nbox.init import nbox_grpc_stub, nbox_ws_v1, nbox_serving_service_stub, nbox_model_service_stub
 from nbox.nbxlib.astea import Astea, IndexTypes as IT
 
@@ -165,7 +166,11 @@ def upload_job_folder(
   init_folder: str,
   id: str = "",
   project_id: str = "",
+
+  # job / deploy rpc things
   trigger: bool = False,
+  deploy: bool = True,
+  pin: bool = False,
 
   # all the things for resources
   resource_cpu: str = "",
@@ -228,10 +233,10 @@ def upload_job_folder(
 
   if method not in OT._valid_deployment_types():
     raise ValueError(f"Invalid method: {method}, should be either {OT._valid_deployment_types()}")
-  # if (not name and not id) or (name and id):
-  #   raise ValueError("Either --name or --id must be present")
-  if trigger and method not in [OT.JOB, OT.SERVING]:
-    raise ValueError(f"Trigger can only be used with '{OT.JOB}' or '{OT.SERVING}'")
+  if trigger and method != OT.JOB:
+    raise ValueError(f"Trigger can only be used with method='{OT.JOB}'")
+  if pin and method != OT.SERVING:
+    raise ValueError(f"Deploy and Pin can only be used with method='{OT.SERVING}'")
   if model_name and method != OT.SERVING:
     raise ValueError(f"model_name can only be used with '{OT.SERVING}'")
 
@@ -410,6 +415,8 @@ def __common_resource(db: Resource) -> Resource:
       },
       exe_jinja_kwargs = exe_jinja_kwargs,
     )
+    if deploy:
+      out.deploy()
     if trigger:
       out.pin()
   else:
@@ -450,7 +457,7 @@ def _get_deployment_data(name: str = "", id: str = "", *, workspace_id: str = ""
   workspace_id = workspace_id or secret(AuthConfig.workspace_id)
 
   # get the deployment
-  serving: Serving = rpc(
+  serving: Serving = mpb.rpc(
     nbox_serving_service_stub.GetServing,
     ServingRequest(
       serving=Serving(name=name, id=id),
@@ -473,7 +480,7 @@ def _get_time(t):
     return datetime.fromtimestamp(int(float(t))).strftime("%Y-%m-%d %H:%M:%S")
 
   workspace_id = workspace_id or secret(AuthConfig.workspace_id)
-  all_deployments: ServingListResponse = rpc(
+  all_deployments: ServingListResponse = mpb.rpc(
     nbox_serving_service_stub.ListServings,
     ServingListRequest(
       auth_info=auth_info_pb(),
@@ -522,6 +529,13 @@ def __init__(self, serving_id: str = "", model_id: str = "", *, workspace_id: st
     self.serving_name = serving_name
     self.ws_stub = nbox_ws_v1.deployments
 
+  def __repr__(self) -> str:
+    x = f"nbox.Serve('{self.id}', '{self.workspace_id}'"
+    if self.model_id is not None:
+      x += f", model_id = '{self.model_id}'"
+    x += ")"
+    return x
+
   def pin(self) -> bool:
     """Pin a model to the deployment
 
@@ -530,7 +544,7 @@ def pin(self) -> bool:
       workspace_id (str, optional): Workspace ID. Defaults to "".
     """
     logger.info(f"Pin model {self.model_id} to deployment {self.serving_id}")
-    rpc(
+    mpb.rpc(
       nbox_model_service_stub.SetModelPin,
       ModelRequest(
         model = ModelProto(
@@ -552,7 +566,7 @@ def unpin(self) -> bool:
       workspace_id (str, optional): Workspace ID. Defaults to "".
     """
     logger.info(f"Unpin model {self.model_id} to deployment {self.serving_id}")
-    rpc(
+    mpb.rpc(
       nbox_model_service_stub.SetModelPin,
       ModelRequest(
         model = ModelProto(
@@ -578,7 +592,7 @@ def scale(self, replicas: int) -> bool:
       raise ValueError("Replicas must be greater than or equal to 0")
 
     logger.info(f"Scale model deployment {self.model_id} to {replicas} replicas")
-    rpc(
+    mpb.rpc(
       nbox_model_service_stub.UpdateModel,
       UpdateModelRequest(
         model=ModelProto(
@@ -600,7 +614,7 @@ def logs(self, f = sys.stdout):
       f (file, optional): File to write the logs to. Defaults to sys.stdout.
     """
     logger.debug(f"Streaming logs of model '{self.model_id}'")
-    for model_log in streaming_rpc(
+    for model_log in mpb.streaming_rpc(
       nbox_model_service_stub.ModelLogs,
       ModelRequest(
         model = ModelProto(
@@ -616,13 +630,24 @@ def logs(self, f = sys.stdout):
         f.write(log)
         f.flush()
 
-  def __repr__(self) -> str:
-    x = f"nbox.Serve('{self.id}', '{self.workspace_id}'"
-    if self.model_id is not None:
-      x += f", model_id = '{self.model_id}'"
-    x += ")"
-    return x
-
+  def deploy(self, tag: str = ""):
+    model = ModelProto(
+      id = self.model_id,
+      serving_group_id = self.serving_id,
+    )
+    if tag:
+      model.feature_gates.update({
+        "SetModelMetadata": tag
+      })
+    response: ModelProto = mpb.rpc(
+      nbox_model_service_stub.Deploy,
+      ModelRequest(
+        model = model,
+        auth_info = auth_info_pb(),
+      ),
+      "Could not deploy model",
+      raise_on_error=True
+    )
 
 
 ################################################################################
@@ -655,7 +680,7 @@ def _get_job_data(name: str = "", id: str = "", remove_archived: bool = True, *,
   if workspace_id == None:
     workspace_id = "personal"
 
-  job: JobProto = rpc(
+  job: JobProto = mpb.rpc(
     nbox_grpc_stub.GetJob,
     JobRequest(
       auth_info = auth_info_pb(),
@@ -683,7 +708,7 @@ def get_job_list(sort: str = "name", *, workspace_id: str = ""):
   def _get_time(t):
     return datetime.fromtimestamp(int(float(t))).strftime("%Y-%m-%d %H:%M:%S")
 
-  out: ListJobsResponse = rpc(
+  out: ListJobsResponse = mpb.rpc(
     nbox_grpc_stub.ListJobs,
     ListJobsRequest(auth_info = auth_info_pb()),
     "Could not get job list",
@@ -794,7 +819,7 @@ def change_schedule(self, new_schedule: Schedule):
     """
     logger.debug(f"Updating job '{self.job_proto.id}'")
     self.job_proto.schedule.MergeFrom(new_schedule.get_message())
-    rpc(
+    mpb.rpc(
       nbox_grpc_stub.UpdateJob,
       UpdateJobRequest(auth_info=self.auth_info, job=self.job_proto, update_mask=FieldMask(paths=["schedule"])),
       "Could not update job schedule",
@@ -814,7 +839,7 @@ def __repr__(self) -> str:
   def logs(self, f = sys.stdout):
     """Stream logs of the job, `f` can be anything has a `.write/.flush` methods"""
     logger.debug(f"Streaming logs of job '{self.job_proto.id}'")
-    for job_log in streaming_rpc(
+    for job_log in mpb.streaming_rpc(
       nbox_grpc_stub.GetJobLogs,
       JobRequest(auth_info=self.auth_info ,job = self.job_proto),
       f"Could not get logs of job {self.job_proto.id}, is your job complete?",
@@ -827,7 +852,7 @@ def logs(self, f = sys.stdout):
   def delete(self):
     """Delete this job"""
     logger.info(f"Deleting job '{self.job_proto.id}'")
-    rpc(nbox_grpc_stub.DeleteJob, JobRequest(auth_info=self.auth_info, job = self.job_proto,), "Could not delete job")
+    mpb.rpc(nbox_grpc_stub.DeleteJob, JobRequest(auth_info=self.auth_info, job = self.job_proto,), "Could not delete job")
     logger.info(f"Deleted job '{self.job_proto.id}'")
     self.refresh()
 
@@ -839,7 +864,7 @@ def refresh(self):
     if self.id == None:
       return
 
-    self.job_proto: JobProto = rpc(
+    self.job_proto: JobProto = mpb.rpc(
       nbox_grpc_stub.GetJob,
       JobRequest(auth_info=self.auth_info, job = self.job_proto),
       f"Could not get job {self.job_proto.id}"
@@ -858,7 +883,7 @@ def trigger(self, tag: str = ""):
     logger.debug(f"Triggering job '{self.job_proto.id}'")
     if tag:
       self.job_proto.feature_gates.update({"SetRunMetadata": tag})
-    rpc(nbox_grpc_stub.TriggerJob, JobRequest(auth_info=self.auth_info, job = self.job_proto), f"Could not trigger job '{self.job_proto.id}'")
+    mpb.rpc(nbox_grpc_stub.TriggerJob, JobRequest(auth_info=self.auth_info, job = self.job_proto), f"Could not trigger job '{self.job_proto.id}'")
     logger.info(f"Triggered job '{self.job_proto.id}'")
     self.refresh()
 
@@ -870,7 +895,7 @@ def pause(self):
     logger.info(f"Pausing job '{self.job_proto.id}'")
     job: JobProto = self.job_proto
     job.status = JobProto.Status.PAUSED
-    rpc(nbox_grpc_stub.UpdateJob, UpdateJobRequest(auth_info=self.auth_info, job=job, update_mask=FieldMask(paths=["status", "paused"])), f"Could not pause job {self.job_proto.id}", True)
+    mpb.rpc(nbox_grpc_stub.UpdateJob, UpdateJobRequest(auth_info=self.auth_info, job=job, update_mask=FieldMask(paths=["status", "paused"])), f"Could not pause job {self.job_proto.id}", True)
     logger.debug(f"Paused job '{self.job_proto.id}'")
     self.refresh()
 
@@ -879,7 +904,7 @@ def resume(self):
     logger.info(f"Resuming job '{self.job_proto.id}'")
     job: JobProto = self.job_proto
     job.status = JobProto.Status.SCHEDULED
-    rpc(nbox_grpc_stub.UpdateJob, UpdateJobRequest(auth_info=self.auth_info, job=job, update_mask=FieldMask(paths=["status", "paused"])), f"Could not resume job {self.job_proto.id}", True)
+    mpb.rpc(nbox_grpc_stub.UpdateJob, UpdateJobRequest(auth_info=self.auth_info, job=job, update_mask=FieldMask(paths=["status", "paused"])), f"Could not resume job {self.job_proto.id}", True)
     logger.debug(f"Resumed job '{self.job_proto.id}'")
     self.refresh()
 

nbox/lib/dist.py

@@ -146,8 +146,8 @@ def serve(self, **serve_kwargs):
         # create the central project class and get the experiment tracker
         proj = Project()
         logger.info(lo("Project data:", **proj.data))
-        live_tracker = proj.get_live_tracker()
-        tracker_config = LiveConfig.from_json(live_tracker.serving.config)
+        # live_tracker = proj.get_live_tracker()
+        # tracker_config = LiveConfig.from_json(live_tracker.serving.config)
 
       # now start serving
       serve_operator(op_or_app = self.op, **serve_kwargs)

nbox/lmao/common.py

@@ -182,21 +182,20 @@ def __init__(
     resource: Resource,
     cli_comm: str,
     enable_system_monitoring: bool = False,
+    extra_kwargs: Dict[str, Any] = {},
   ):
     self.resource = resource
     self.cli_comm = cli_comm
     self.enable_system_monitoring = enable_system_monitoring
-    self.keys = set()
+    self.extra_kwargs = extra_kwargs
 
   def to_dict(self):
-    out = {
+    return {
       "resource": mpb.message_to_dict(self.resource),
       "cli_comm": self.cli_comm,
       "enable_system_monitoring": self.enable_system_monitoring,
+      "extra_kwargs": self.extra_kwargs,
     }
-    for k in self.keys:
-      out[k] = getattr(self, k)
-    return out
 
   def to_json(self):
     return json.dumps(self.to_dict())
@@ -206,16 +205,7 @@ def from_json(cls, json_str) -> 'LiveConfig':
     d = json.loads(json_str)
     d["resource"] = resource_from_dict(d["resource"])
     _cls = cls(**d)
-    for k in d:
-      if k not in ["resource", "cli_comm", "enable_system_monitoring"]:
-        _cls.add(k, d[k])
-  
-  def add(self, key, value):
-    setattr(self, key, value)
-    self.keys.add(key)
-
-  def get(self, key):
-    return getattr(self, key)
+    return _cls
 
 
 """

nbox/lmao/exp.py

@@ -16,7 +16,7 @@
 from nbox.lmao.lmao_rpc_client import (
   AgentDetails,
   RunLog,
-  Run,
+  Run as RunProto,
   InitRunRequest
 )
 from nbox.observability.system import SystemMetricsLogger
@@ -111,7 +111,7 @@ def __init__(
       raise Exception(f"Project with id {self.project_id} does not exist")
 
     self.config = self._get_config(metadata = metadata)
-    self.run = self._init_experiment(
+    self.run: RunProto = self._init_experiment(
       project_id = project_id,
       config = self.config,
       experiment_id = experiment_id
@@ -148,7 +148,7 @@ def _get_config(self, metadata: Dict[str, Any]):
       log_config["git"] = get_git_details("./")
     return log_config
 
-  def _init_experiment(self, project_id, config: Dict[str, Any] = {}, experiment_id: str = ""):
+  def _init_experiment(self, project_id, config: Dict[str, Any] = {}, experiment_id: str = "") -> RunProto:
     # update the server or create new experiment
     agent_details = AgentDetails(
       workspace_id = self.workspace_id,
@@ -163,7 +163,7 @@ def _init_experiment(self, project_id, config: Dict[str, Any] = {}, experiment_i
 
     if experiment_id:
       action = "Updated"
-      run_details = self.lmao.get_run_details(Run(
+      run_details = self.lmao.get_run_details(RunProto(
         workspace_id = self.workspace_id,
         project_id = project_id,
         experiment_id = experiment_id,
@@ -173,7 +173,7 @@ def _init_experiment(self, project_id, config: Dict[str, Any] = {}, experiment_i
         raise Exception("Server Side exception has occurred, Check the log for details")
       if run_details.experiment_id:
         # means that this run already exists so we need to make an update call
-        ack = self.lmao.update_run_status(Run(
+        ack = self.lmao.update_run_status(RunProto(
           workspace_id = self.workspace_id,
           project_id = project_id,
           experiment_id = run_details.experiment_id,
@@ -204,6 +204,10 @@ def _init_experiment(self, project_id, config: Dict[str, Any] = {}, experiment_i
 
   """The functions below are the ones supposed to be used."""
 
+  @property
+  def run_config(self) -> Dict[str, Any]:
+    return loads(self.run.config)
+
   @lru_cache(maxsize=1)
   def get_relic(self):
     """Get the underlying Relic for more advanced usage patterns."""
@@ -285,3 +289,4 @@ def add_files(self, *files: List[str]):
       relic = self.get_relic()
       for f in all_files:
         relic.put(f)
+    return all_files