trainml 0.5.17__tar.gz → 1.0.0__tar.gz

{trainml-0.5.17 → trainml-1.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: trainml
-Version: 0.5.17
+Version: 1.0.0
 Summary: trainML client SDK and command line utilities
 Home-page: https://github.com/trainML/trainml-cli
 Author: trainML
@@ -177,7 +177,7 @@ Description: <div align="center">
         trainml dataset list
         ```
         
-        To connect to a job that requires the [connection capability](https://docs.trainml.ai/reference/connection-capability):
+        To connect to a job that uses the "local" file transfer method:
         
         ```
         trainml job connect <job ID or name>

{trainml-0.5.17 → trainml-1.0.0}/README.md

@@ -169,7 +169,7 @@ To list all datasets:
 trainml dataset list
 ```
 
-To connect to a job that requires the [connection capability](https://docs.trainml.ai/reference/connection-capability):
+To connect to a job that uses the "local" file transfer method:
 
 ```
 trainml job connect <job ID or name>

{trainml-0.5.17 → trainml-1.0.0}/examples/local_storage.py

@@ -19,7 +19,6 @@ async def create_dataset():
     attach_task = asyncio.create_task(dataset.attach())
     connect_task = asyncio.create_task(dataset.connect())
     await asyncio.gather(attach_task, connect_task)
-    await dataset.disconnect()
     return dataset
 
 
@@ -55,7 +54,6 @@ async def run_job(dataset):
     await asyncio.gather(attach_task, connect_task)
 
     # Cleanup job
-    await job.disconnect()
     await job.remove()
 
 

{trainml-0.5.17 → trainml-1.0.0}/pyproject.toml

@@ -25,6 +25,7 @@ markers = [
     "data_connectors: Data Connector tests",
     "services: Services tests",
     "device_configs: DeviceConfigs tests",
+    "local: Local Connection Utility tests",
     "unit: All unit tests (no trainML environment required)",
     "integration: All integration tests (trainML environment required)",
     "sdk: All tests of the SDK",

{trainml-0.5.17 → trainml-1.0.0}/tests/integration/test_checkpoints_integration.py

@@ -23,6 +23,7 @@ class GetCheckpointTests:
         checkpoint = await checkpoint.wait_for("archived", 60)
 
     async def test_get_checkpoints(self, trainml, checkpoint):
+        _ = checkpoint
         checkpoints = await trainml.checkpoints.list()
         assert len(checkpoints) > 0
 
@@ -55,7 +56,7 @@ class GetCheckpointTests:
 
 @mark.create
 @mark.asyncio
-async def test_checkpoint_wasabi(trainml, capsys):
+async def test_checkpoint_wasabi(trainml):
     checkpoint = await trainml.checkpoints.create(
         name="CLI Automated Wasabi",
         source_type="wasabi",
@@ -72,6 +73,7 @@ async def test_checkpoint_wasabi(trainml, capsys):
 
 @mark.create
 @mark.asyncio
+@mark.local
 async def test_checkpoint_local(trainml, capsys):
     checkpoint = await trainml.checkpoints.create(
         name="CLI Automated Local",
@@ -81,7 +83,6 @@ async def test_checkpoint_local(trainml, capsys):
     attach_task = asyncio.create_task(checkpoint.attach())
     connect_task = asyncio.create_task(checkpoint.connect())
     await asyncio.gather(attach_task, connect_task)
-    await checkpoint.disconnect()
     await checkpoint.refresh()
     status = checkpoint.status
     size = checkpoint.size
@@ -92,5 +93,5 @@ async def test_checkpoint_local(trainml, capsys):
     sys.stdout.write(captured.out)
     sys.stderr.write(captured.err)
     assert "Starting data upload from local" in captured.out
-    assert "official/LICENSE  11456 bytes" in captured.out
+    assert "official/LICENSE" in captured.out
     assert "Upload complete" in captured.out

{trainml-0.5.17 → trainml-1.0.0}/tests/integration/test_datasets_integration.py

@@ -50,7 +50,9 @@ class GetDatasetTests:
     async def test_dataset_repr(self, dataset):
         string = repr(dataset)
         regex = (
-            r"^Dataset\( trainml , \*\*{.*'dataset_uuid': '" + dataset.id + r"'.*}\)$"
+            r"^Dataset\( trainml , \*\*{.*'dataset_uuid': '"
+            + dataset.id
+            + r"'.*}\)$"
         )
         assert isinstance(string, str)
         assert re.match(regex, string)
@@ -79,6 +81,7 @@ class GetDatasetTests:
 
 @mark.create
 @mark.asyncio
+@mark.local
 async def test_dataset_local(trainml, capsys):
     dataset = await trainml.datasets.create(
         name="CLI Automated Local",
@@ -88,7 +91,6 @@ async def test_dataset_local(trainml, capsys):
     attach_task = asyncio.create_task(dataset.attach())
     connect_task = asyncio.create_task(dataset.connect())
     await asyncio.gather(attach_task, connect_task)
-    await dataset.disconnect()
     await dataset.refresh()
     status = dataset.status
     size = dataset.size
@@ -99,5 +101,5 @@ async def test_dataset_local(trainml, capsys):
     sys.stdout.write(captured.out)
     sys.stderr.write(captured.err)
     assert "Starting data upload from local" in captured.out
-    assert "data_batch_1.bin  30733788 bytes" in captured.out
+    assert "data_batch_1.bin" in captured.out
     assert "Upload complete" in captured.out

{trainml-0.5.17 → trainml-1.0.0}/tests/integration/test_jobs_integration.py

@@ -46,7 +46,10 @@ class JobLifeCycleTests:
         job = await job.wait_for("running")
         assert job.status == "running"
         assert job.url
-        assert extract_domain_suffix(urlparse(job.url).hostname) == "proximl.cloud"
+        assert (
+            extract_domain_suffix(urlparse(job.url).hostname)
+            == "proximl.cloud"
+        )
 
     async def test_stop_job(self, job):
         assert job.status == "running"
@@ -204,7 +207,8 @@ class JobAPIResourceValidationTests:
                 disk_size=10,
             )
         assert (
-            "Invalid Request - CPU Count must be a multiple of 4" in error.value.message
+            "Invalid Request - CPU Count must be a multiple of 4"
+            in error.value.message
         )
 
     async def test_invalid_gpu_count_for_cpu(self, trainml):
@@ -417,6 +421,7 @@ class JobAPIWorkerValidationTests:
 @mark.asyncio
 @mark.xdist_group("job_io")
 class JobIOTests:
+    @mark.local
     async def test_job_local_output(self, trainml, capsys):
         temp_dir = tempfile.TemporaryDirectory()
         job = await trainml.jobs.create(
@@ -426,7 +431,7 @@ class JobIOTests:
             disk_size=10,
             workers=["python $ML_MODEL_PATH/tensorflow/main.py"],
             environment=dict(
-                type="DEEPLEARNING_PY312",
+                type="DEEPLEARNING_PY313",
                 env=[
                     dict(
                         key="CHECKPOINT_FILE",
@@ -452,13 +457,13 @@ class JobIOTests:
                 ],
             ),
         )
-        await job.wait_for("waiting for data/model download")
+        # Wait for job to reach running status since only output_type is local
+        await job.wait_for("running")
         attach_task = asyncio.create_task(job.attach())
         connect_task = asyncio.create_task(job.connect())
         await asyncio.gather(attach_task, connect_task)
         await job.refresh()
         assert job.status == "finished"
-        await job.disconnect()
         await job.remove()
         upload_contents = os.listdir(temp_dir.name)
         temp_dir.cleanup()
@@ -470,9 +475,8 @@ class JobIOTests:
         captured = capsys.readouterr()
         sys.stdout.write(captured.out)
         sys.stderr.write(captured.err)
-        assert "Epoch 1/2" in captured.out
-        assert "Epoch 2/2" in captured.out
-        assert "adding: model.ckpt-0001" in captured.out
+        assert "Epoch 1/2" in captured.out or "Epoch 2/2" in captured.out
+        assert "model.ckpt-0001" in captured.out
         assert "Send complete" in captured.out
 
     async def test_job_model_input_and_output(self, trainml, capsys):
@@ -513,8 +517,7 @@ class JobIOTests:
         captured = capsys.readouterr()
         sys.stdout.write(captured.out)
         sys.stderr.write(captured.err)
-        assert "Epoch 1/2" in captured.out
-        assert "Epoch 2/2" in captured.out
+        assert "Epoch 1/2" in captured.out or "Epoch 2/2" in captured.out
 
         new_model = await trainml.models.get(workers[0].get("output_uuid"))
         assert new_model.id
@@ -560,9 +563,12 @@ class JobTypeTests:
         await job.wait_for("running")
         await job.refresh()
         assert job.url
-        assert extract_domain_suffix(urlparse(job.url).hostname) == "proximl.cloud"
+        assert (
+            extract_domain_suffix(urlparse(job.url).hostname)
+            == "proximl.cloud"
+        )
         tries = 0
-        await asyncio.sleep(180) ## downloading weights can be slow
+        await asyncio.sleep(180)  ## downloading weights can be slow
         async with aiohttp.ClientSession() as session:
             retry = True
             while retry:
@@ -640,9 +646,11 @@ class JobTypeTests:
         captured = capsys.readouterr()
         sys.stdout.write(captured.out)
         sys.stderr.write(captured.err)
-        assert "Epoch 1/2" in captured.out
-        assert "Epoch 2/2" in captured.out
-        assert "Uploading s3://trainml-example/output/resnet_cifar10" in captured.out
+        assert "Epoch 1/2" in captured.out or "Epoch 2/2" in captured.out
+        assert (
+            "Uploading s3://trainml-example/output/resnet_cifar10"
+            in captured.out
+        )
         assert (
             "upload: ./model.ckpt-0002.weights.h5 to s3://trainml-example/output/resnet_cifar10/model.ckpt-0002.weights.h5"
             in captured.out
@@ -680,9 +688,12 @@ class JobFeatureTests:
         captured = capsys.readouterr()
         sys.stdout.write(captured.out)
         sys.stderr.write(captured.err)
-        assert "Train Epoch: 1 [0/60000 (0%)]" in captured.out
-        assert "Train Epoch: 1 [59520/60000 (99%)]" in captured.out
+        assert (
+            "Train Epoch: 1 [0/60000 (0%)]" in captured.out
+            or "Train Epoch: 1 [59520/60000 (99%)]" in captured.out
+        )
 
+    @mark.local
     async def test_inference_job(self, trainml, capsys):
         temp_dir = tempfile.TemporaryDirectory()
         job = await trainml.jobs.create(
@@ -706,11 +717,11 @@ class JobFeatureTests:
         )
         assert job.id
         await job.wait_for("running")
-        await job.connect()
-        await job.attach()
+        attach_task = asyncio.create_task(job.attach())
+        connect_task = asyncio.create_task(job.connect())
+        await asyncio.gather(attach_task, connect_task)
         await job.refresh()
         assert job.status == "finished"
-        await job.disconnect()
         await job.remove()
         await job.wait_for("archived")
         captured = capsys.readouterr()
@@ -719,15 +730,10 @@ class JobFeatureTests:
         upload_contents = os.listdir(temp_dir.name)
         temp_dir.cleanup()
         assert len(upload_contents) >= 3
-        assert any(
-            "model.ckpt-0002" in content
-            for content in upload_contents
-        )
+        assert any("model.ckpt-0002" in content for content in upload_contents)
 
         captured = capsys.readouterr()
         sys.stdout.write(captured.out)
         sys.stderr.write(captured.err)
-        assert "Epoch 1/2" in captured.out
-        assert "Epoch 2/2" in captured.out
-        assert "Number of regular files transferred: 4" in captured.out
+        assert "Epoch 1/2" in captured.out or "Epoch 2/2" in captured.out
         assert "Send complete" in captured.out

{trainml-0.5.17 → trainml-1.0.0}/tests/integration/test_models_integration.py

@@ -44,7 +44,11 @@ class GetModelTests:
 
     async def test_model_repr(self, model):
         string = repr(model)
-        regex = r"^Model\( trainml , \*\*{.*'model_uuid': '" + model.id + r"'.*}\)$"
+        regex = (
+            r"^Model\( trainml , \*\*{.*'model_uuid': '"
+            + model.id
+            + r"'.*}\)$"
+        )
         assert isinstance(string, str)
         assert re.match(regex, string)
 
@@ -68,6 +72,7 @@ async def test_model_wasabi(trainml, capsys):
 
 @mark.create
 @mark.asyncio
+@mark.local
 async def test_model_local(trainml, capsys):
     model = await trainml.models.create(
         name="CLI Automated Local",
@@ -77,7 +82,6 @@ async def test_model_local(trainml, capsys):
     attach_task = asyncio.create_task(model.attach())
     connect_task = asyncio.create_task(model.connect())
     await asyncio.gather(attach_task, connect_task)
-    await model.disconnect()
     await model.refresh()
     status = model.status
     size = model.size
@@ -88,5 +92,5 @@ async def test_model_local(trainml, capsys):
     sys.stdout.write(captured.out)
     sys.stderr.write(captured.err)
     assert "Starting data upload from local" in captured.out
-    assert "official/LICENSE  11456 bytes" in captured.out
+    assert "official/LICENSE" in captured.out
     assert "Upload complete" in captured.out

{trainml-0.5.17 → trainml-1.0.0}/tests/integration/test_volumes_integration.py

@@ -74,6 +74,7 @@ async def test_volume_wasabi(trainml, capsys):
 
 @mark.create
 @mark.asyncio
+@mark.local
 async def test_volume_local(trainml, capsys):
     volume = await trainml.volumes.create(
         name="CLI Automated Local",
@@ -84,7 +85,6 @@ async def test_volume_local(trainml, capsys):
     attach_task = asyncio.create_task(volume.attach())
     connect_task = asyncio.create_task(volume.connect())
     await asyncio.gather(attach_task, connect_task)
-    await volume.disconnect()
     await volume.refresh()
     status = volume.status
     billed_size = volume.billed_size
@@ -97,5 +97,5 @@ async def test_volume_local(trainml, capsys):
     sys.stdout.write(captured.out)
     sys.stderr.write(captured.err)
     assert "Starting data upload from local" in captured.out
-    assert "official/LICENSE  11456 bytes" in captured.out
+    assert "official/LICENSE" in captured.out
     assert "Upload complete" in captured.out

trainml-1.0.0/tests/unit/cli/test_cli_checkpoint_unit.py

@@ -0,0 +1,333 @@
+import re
+import json
+import click
+from unittest.mock import AsyncMock, patch, Mock
+from pytest import mark, fixture, raises
+
+pytestmark = [mark.cli, mark.unit, mark.checkpoints]
+
+from trainml.cli import checkpoint as specimen
+from trainml.cli.checkpoint import pretty_size
+from trainml.checkpoints import Checkpoint
+
+
+def test_pretty_size_zero():
+    """Test pretty_size with zero/None (line 7)."""
+    result = pretty_size(None)
+    assert result == "0.00   B"
+    result = pretty_size(0)
+    assert result == "0.00   B"
+
+
+def test_list(runner, mock_my_checkpoints):
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+        mock_trainml.checkpoints = AsyncMock()
+        mock_trainml.checkpoints.list = AsyncMock(
+            return_value=mock_my_checkpoints
+        )
+        result = runner.invoke(specimen, ["list"])
+        print(result)
+        assert result.exit_code == 0
+        mock_trainml.checkpoints.list.assert_called_once()
+
+
+def test_attach_success(runner, mock_my_checkpoints):
+    """Test attach command success (lines 32-38)."""
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+
+        async def list_async():
+            return mock_my_checkpoints
+
+        mock_trainml.checkpoints = AsyncMock()
+        mock_trainml.checkpoints.list = Mock(side_effect=lambda: list_async())
+
+        # Use the first checkpoint from the list
+        checkpoint = mock_my_checkpoints[0]
+
+        async def attach_async():
+            return None
+
+        checkpoint.attach = Mock(return_value=attach_async())
+
+        with patch("trainml.cli.search_by_id_name", return_value=checkpoint):
+            result = runner.invoke(specimen, ["attach", "1"])
+            assert result.exit_code == 0
+            checkpoint.attach.assert_called_once()
+
+
+def test_attach_not_found(runner, mock_my_checkpoints):
+    """Test attach command when checkpoint not found (line 36)."""
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+
+        async def list_async():
+            return mock_my_checkpoints
+
+        mock_trainml.checkpoints = AsyncMock()
+        mock_trainml.checkpoints.list = Mock(side_effect=lambda: list_async())
+
+        with patch("trainml.cli.search_by_id_name", return_value=None):
+            result = runner.invoke(specimen, ["attach", "nonexistent"])
+            assert result.exit_code != 0
+            assert "Cannot find specified checkpoint" in result.output
+
+
+def test_connect_with_attach(runner, mock_my_checkpoints):
+    """Test connect command with attach (lines 56-65, attach=True)."""
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+
+        async def list_async():
+            return mock_my_checkpoints
+
+        mock_trainml.checkpoints = AsyncMock()
+        mock_trainml.checkpoints.list = Mock(side_effect=lambda: list_async())
+
+        checkpoint = mock_my_checkpoints[0]
+
+        async def connect_async():
+            return None
+
+        async def attach_async():
+            return None
+
+        checkpoint.connect = Mock(return_value=connect_async())
+        checkpoint.attach = Mock(return_value=attach_async())
+
+        with patch("trainml.cli.search_by_id_name", return_value=checkpoint):
+            result = runner.invoke(specimen, ["connect", "1"])
+            assert result.exit_code == 0
+            checkpoint.connect.assert_called_once()
+            checkpoint.attach.assert_called_once()
+
+
+def test_connect_no_attach(runner, mock_my_checkpoints):
+    """Test connect command without attach (lines 56-65, attach=False)."""
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+
+        async def list_async():
+            return mock_my_checkpoints
+
+        mock_trainml.checkpoints = AsyncMock()
+        mock_trainml.checkpoints.list = Mock(side_effect=lambda: list_async())
+
+        checkpoint = mock_my_checkpoints[0]
+
+        async def connect_async():
+            return None
+
+        checkpoint.connect = Mock(return_value=connect_async())
+
+        with patch("trainml.cli.search_by_id_name", return_value=checkpoint):
+            result = runner.invoke(specimen, ["connect", "--no-attach", "1"])
+            assert result.exit_code == 0
+            checkpoint.connect.assert_called_once()
+
+
+def test_connect_not_found(runner, mock_my_checkpoints):
+    """Test connect command when checkpoint not found (line 60)."""
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+
+        async def list_async():
+            return mock_my_checkpoints
+
+        mock_trainml.checkpoints = AsyncMock()
+        mock_trainml.checkpoints.list = Mock(side_effect=lambda: list_async())
+
+        with patch("trainml.cli.search_by_id_name", return_value=None):
+            result = runner.invoke(specimen, ["connect", "nonexistent"])
+            assert result.exit_code != 0
+            assert "Cannot find specified checkpoint" in result.output
+
+
+def test_create_with_connect_and_attach(runner, tmp_path, mock_my_checkpoints):
+    """Test create command with connect and attach (lines 103-115)."""
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+        checkpoint = mock_my_checkpoints[0]
+
+        async def connect_async():
+            return None
+
+        async def attach_async():
+            return None
+
+        checkpoint.connect = Mock(return_value=connect_async())
+        checkpoint.attach = Mock(return_value=attach_async())
+
+        async def create_async(**kwargs):
+            return checkpoint
+
+        mock_trainml.checkpoints = AsyncMock()
+        mock_trainml.checkpoints.create = Mock(
+            side_effect=lambda **kwargs: create_async(**kwargs)
+        )
+
+        test_dir = tmp_path / "test_checkpoint"
+        test_dir.mkdir()
+        result = runner.invoke(
+            specimen, ["create", "test-checkpoint", str(test_dir)]
+        )
+        assert result.exit_code == 0
+        mock_trainml.checkpoints.create.assert_called_once()
+        checkpoint.connect.assert_called_once()
+        checkpoint.attach.assert_called_once()
+
+
+def test_create_with_connect_no_attach(runner, tmp_path, mock_my_checkpoints):
+    """Test create command with connect but no attach (lines 103-115)."""
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+        checkpoint = mock_my_checkpoints[0]
+
+        async def connect_async():
+            return None
+
+        checkpoint.connect = Mock(return_value=connect_async())
+
+        async def create_async(**kwargs):
+            return checkpoint
+
+        mock_trainml.checkpoints = AsyncMock()
+        mock_trainml.checkpoints.create = Mock(
+            side_effect=lambda **kwargs: create_async(**kwargs)
+        )
+
+        test_dir = tmp_path / "test_checkpoint"
+        test_dir.mkdir()
+        result = runner.invoke(
+            specimen,
+            ["create", "--no-attach", "test-checkpoint", str(test_dir)],
+        )
+        assert result.exit_code == 0
+        checkpoint.connect.assert_called_once()
+
+
+def test_create_no_connect(runner, tmp_path):
+    """Test create command without connect (lines 103-115, line 115)."""
+    mock_checkpoint = Mock(spec=Checkpoint)
+
+    mock_trainml_runner = Mock()
+    mock_trainml_runner.client = Mock()
+    mock_trainml_runner.client.checkpoints = Mock()
+    mock_trainml_runner.client.checkpoints.create = AsyncMock(
+        return_value=mock_checkpoint
+    )
+    mock_trainml_runner.run = Mock(
+        side_effect=lambda x: x if not hasattr(x, "__call__") else x()
+    )
+
+    with patch("trainml.cli.TrainMLRunner", return_value=mock_trainml_runner):
+        test_dir = tmp_path / "test_checkpoint"
+        test_dir.mkdir()
+        result = runner.invoke(
+            specimen,
+            ["create", "--no-connect", "test-checkpoint", str(test_dir)],
+        )
+        assert result.exit_code != 0
+        assert "No logs to show" in result.output
+
+
+def test_list_public(runner, mock_my_checkpoints):
+    """Test list_public command (lines 152-171)."""
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+        mock_trainml.checkpoints = AsyncMock()
+        mock_trainml.checkpoints.list_public = AsyncMock(
+            return_value=mock_my_checkpoints
+        )
+
+        result = runner.invoke(specimen, ["list-public"])
+        assert result.exit_code == 0
+        mock_trainml.checkpoints.list_public.assert_called_once()
+
+
+def test_remove_success(runner, mock_my_checkpoints):
+    """Test remove command success (lines 192-201)."""
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+
+        async def list_async():
+            return mock_my_checkpoints
+
+        mock_trainml.checkpoints = AsyncMock()
+        mock_trainml.checkpoints.list = Mock(side_effect=lambda: list_async())
+
+        checkpoint = mock_my_checkpoints[0]
+
+        async def remove_async():
+            return None
+
+        checkpoint.remove = Mock(return_value=remove_async())
+
+        with patch("trainml.cli.search_by_id_name", return_value=checkpoint):
+            result = runner.invoke(specimen, ["remove", "1"])
+            assert result.exit_code == 0
+            checkpoint.remove.assert_called_once_with(force=False)
+
+
+def test_remove_not_found(runner, mock_my_checkpoints):
+    """Test remove command when checkpoint not found (lines 192-201)."""
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+
+        async def list_async():
+            return mock_my_checkpoints
+
+        mock_trainml.checkpoints = AsyncMock()
+        mock_trainml.checkpoints.list = Mock(side_effect=lambda: list_async())
+
+        with patch("trainml.cli.search_by_id_name", return_value=None):
+            result = runner.invoke(specimen, ["remove", "nonexistent"])
+            assert result.exit_code != 0
+            assert "Cannot find specified checkpoint" in result.output
+
+
+def test_rename_success(runner, mock_my_checkpoints):
+    """Test rename command success (lines 214-223)."""
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+        checkpoint = mock_my_checkpoints[0]
+
+        async def rename_async():
+            return None
+
+        checkpoint.rename = Mock(return_value=rename_async())
+
+        async def get_async(checkpoint_id):
+            return checkpoint
+
+        mock_trainml.checkpoints = AsyncMock()
+        mock_trainml.checkpoints.get = Mock(
+            side_effect=lambda checkpoint_id: get_async(checkpoint_id)
+        )
+
+        result = runner.invoke(specimen, ["rename", "1", "new-name"])
+        assert result.exit_code == 0
+        checkpoint.rename.assert_called_once_with(name="new-name")
+
+
+def test_rename_not_found_none(runner):
+    """Test rename command when checkpoint is None (lines 214-223, line 219)."""
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+
+        async def get_async(checkpoint_id):
+            return None
+
+        mock_trainml.checkpoints = AsyncMock()
+        mock_trainml.checkpoints.get = Mock(
+            side_effect=lambda checkpoint_id: get_async(checkpoint_id)
+        )
+
+        result = runner.invoke(specimen, ["rename", "nonexistent", "new-name"])
+        assert result.exit_code != 0
+        assert "Cannot find specified checkpoint" in result.output
+
+
+def test_rename_not_found_exception(runner):
+    """Test rename command when exception occurs (lines 214-223, line 221)."""
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+
+        async def get_async(checkpoint_id):
+            raise Exception("Not found")
+
+        mock_trainml.checkpoints = AsyncMock()
+        mock_trainml.checkpoints.get = Mock(
+            side_effect=lambda checkpoint_id: get_async(checkpoint_id)
+        )
+
+        result = runner.invoke(specimen, ["rename", "nonexistent", "new-name"])
+        assert result.exit_code != 0
+        assert "Cannot find specified checkpoint" in result.output