feature gates in jobs + deploy

Author: yashbonde

nbox/jobs.py

@@ -171,6 +171,7 @@ def upload_job_folder(
   trigger: bool = False,
   deploy: bool = True,
   pin: bool = False,
+  feature_gates: dict = {},
 
   # all the things for resources
   resource_cpu: str = "",
@@ -380,6 +381,7 @@ def __common_resource(db: Resource) -> Resource:
       init_folder = init_folder,
       job_id = job_proto.id,
       job_name = job_proto.name,
+      feature_gates = feature_gates,
       dag = DAGProto(),
       workspace_id = workspace_id,
       schedule = None,
@@ -416,7 +418,7 @@ def __common_resource(db: Resource) -> Resource:
       exe_jinja_kwargs = exe_jinja_kwargs,
     )
     if deploy:
-      out.deploy()
+      out.deploy(feature_gates = feature_gates)
     if trigger:
       out.pin()
   else:
@@ -630,14 +632,15 @@ def logs(self, f = sys.stdout):
         f.write(log)
         f.flush()
 
-  def deploy(self, tag: str = ""):
+  def deploy(self, tag: str = "", feature_gates: Dict[str, str] = {}):
     model = ModelProto(
       id = self.model_id,
       serving_group_id = self.serving_id,
     )
     if tag:
       model.feature_gates.update({
-        "SetModelMetadata": tag
+        "SetModelMetadata": tag,
+        **feature_gates,
       })
     response: ModelProto = mpb.rpc(
       nbox_model_service_stub.Deploy,
@@ -874,14 +877,15 @@ def refresh(self):
 
     self.status = self.job_proto.Status.keys()[self.job_proto.status]
 
-  def trigger(self, tag: str = ""):
+  def trigger(self, tag: str = "", feature_gates: dict = {}):
     """Manually triger this job.
     
     Args:
       tag (str, optional): Tag to be set in the run metadata, read in more detail before trying to use this. Defaults to "".
     """
     logger.debug(f"Triggering job '{self.job_proto.id}'")
     if tag:
+      self.job_proto.feature_gates.update(feature_gates)
       self.job_proto.feature_gates.update({"SetRunMetadata": tag})
     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}'")

nbox/network.py

@@ -209,6 +209,7 @@ def _upload_serving_zip(zip_path: str, workspace_id: str, serving_id: str, model
 def deploy_job(
   init_folder: str,
   job_name: str,
+  feature_gates: Dict[str, str],
   dag: DAG,
   schedule: Schedule,
   resource: Resource,
@@ -256,7 +257,8 @@ def deploy_job(
     created_at = SimplerTimes.get_now_pb(),
     schedule = schedule.get_message() if schedule is not None else None, # JobProto.Schedule(cron = "0 0 24 2 0"),
     dag = dag,
-    resource = resource
+    resource = resource,
+    feature_gates = feature_gates,
   )
   job_proto_fp = U.join(gettempdir(), "job_proto.msg")
   mpb.write_binary_to_file(job_proto, job_proto_fp)

nbox/relics/client.py

@@ -263,9 +263,8 @@ def get_from(self, local_path: str, remote_path: str, unzip: bool = False) -> No
       raise ValueError("Relic does not exist, pass create=True")
     logger.debug(f"Getting '{local_path}' from '{remote_path}'")
 
-    files = self.list_files(remote_path)
     file_ = list(filter(
-      lambda x: x.name == remote_path, files
+      lambda x: x.name == remote_path, self.ls(remote_path)
     ))
     if not file_:
       raise ValueError(f"File {remote_path} does not exist in the relic")

nbox/version.py

@@ -1,4 +1,4 @@
-__version__ = "0.12.5c9"
+__version__ = "0.12.5c10"
 _major, _minor, _patch = __version__.split(".")
 _major = int(_major)
 _minor = int(_minor)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nbox"
-version = "0.12.5c9"
+version = "0.12.5c10"
 description = "ML Inference 🥶"
 authors = [
   "NBX Research <research@nimblebox.ai>",