Merge branch 'master' into hoxo/serialize-mount-operations

Author: imakunin

lzy/allocator/src/main/java/ai/lzy/allocator/AllocatorMain.java

@@ -178,8 +178,6 @@ public void destroyAllForTests() throws SQLException {
                 }
             }
         });
-
-        System.out.println("!!!");
     }
 
     public static void main(String[] args) throws IOException, InterruptedException {

lzy/allocator/src/main/java/ai/lzy/allocator/services/AllocatorService.java

@@ -589,8 +589,8 @@ public void forceFree(ForceFreeRequest request, StreamObserver<LongRunning.Opera
                                     ForceFreeResponse.getDefaultInstance()));
                         }
 
-                        if (vm.status() != Vm.Status.RUNNING) {
-                            LOG.error("ForceFree vm {} in status {}, expected RUNNING", vm, vm.status());
+                        if (vm.status() != Vm.Status.RUNNING && vm.status() != Vm.Status.IDLE) {
+                            LOG.error("ForceFree vm {} in status {}, expected RUNNING or IDLE", vm, vm.status());
                             return Either.left(Status.FAILED_PRECONDITION.withDescription("State is " + vm.status()));
                         }
 
@@ -616,7 +616,11 @@ public void forceFree(ForceFreeRequest request, StreamObserver<LongRunning.Opera
                         LOG.info("VM {} scheduled to remove (force free)", vm.vmId());
 
                         allocationContext.startNew(action);
-                        allocationContext.metrics().runningVms.labels(vm.poolLabel()).dec();
+
+                        switch (vm.status()) {
+                            case RUNNING -> allocationContext.metrics().runningVms.labels(vm.poolLabel()).dec();
+                            case IDLE -> allocationContext.metrics().cachedVms.labels(vm.poolLabel()).dec();
+                        }
 
                         return Either.right(op);
                     }

lzy/allocator/src/test/java/ai/lzy/allocator/test/AllocatorServiceTest.java

@@ -438,6 +438,26 @@ public void allocateFromCache() throws Exception {
         assertVmMetrics("S", 0, 0, 0);
     }
 
+    @Test
+    public void forceFreeFromCache() throws Exception {
+        var sessionId = createSession(Durations.fromSeconds(10));
+
+        final CountDownLatch kuberRemoveRequestLatch = new CountDownLatch(1);
+
+        var firstVm = allocateAndFreeVm(sessionId,
+            vm -> mockDeletePodByName(vm.podName(), kuberRemoveRequestLatch::countDown, HttpURLConnection.HTTP_OK));
+
+        assertVmMetrics("S", -1, 0, 1);
+
+        var forceFreeOp = forceFreeVm(firstVm.vmId());
+        waitOpSuccess(forceFreeOp);
+        assertVmMetrics("S", -1, 0, 0);
+
+        var secondVm = allocateVm(sessionId, "S", null);
+        Assert.assertNotEquals(firstVm.vmId(), secondVm.vmId());
+        assertVmMetrics("S", -1, 1, 0);
+    }
+
     @Test
     public void allocateConcurrentFromCache() throws Exception {
         var sessionId = createSession(Durations.fromSeconds(5000));

pylzy/lzy/api/v1/env.py

@@ -12,10 +12,8 @@
 
 import yaml
 
-from pypi_simple import PYPI_SIMPLE_ENDPOINT
-
-from lzy.py_env.api import PyEnv
 from lzy.logs.config import get_logger
+from lzy.py_env.api import PyEnv
 from lzy.utils.pypi import validate_pypi_index_url
 
 _LOG = get_logger(__name__)
@@ -158,7 +156,7 @@ def generate_conda_config(self) -> Dict[str, Any]:
         if python_version in _INSTALLED_VERSIONS:
             env_name = _INSTALLED_VERSIONS[python_version]
         else:
-            _LOG.warn(
+            _LOG.warning(
                 f"Installed python version ({python_version}) is not cached remotely. "
                 f"Usage of a cached python version ({list(_INSTALLED_VERSIONS.keys())}) "
                 f"can decrease startup time."

pylzy/lzy/api/v1/utils/proxy_adapter.py

@@ -45,7 +45,7 @@ def materialize_if_sequence_of_lzy_proxies(obj: Any) -> Any:
         if is_lzy_proxy(elem):
             result.append(materialize(elem))
 
-    if type(obj) == tuple:
+    if type(obj) is tuple:
         return tuple(result)
     return result
 

pylzy/tests/api/v1/test_workflow.py

@@ -356,7 +356,7 @@ def test_materialize(self):
     def test_lazy_args_loading(self):
         @op(lazy_arguments=True)
         def is_arg_type_str(a: str) -> bool:
-            return type(a) == str
+            return type(a) is str
 
         with self.lzy.workflow("test"):
             res = is_arg_type_str("str")
@@ -365,7 +365,7 @@ def is_arg_type_str(a: str) -> bool:
     def test_eager_args_loading(self):
         @op(lazy_arguments=False)
         def is_arg_type_str(a: str) -> bool:
-            return type(a) == str
+            return type(a) is str
 
         with self.lzy.workflow("test"):
             res = is_arg_type_str("str")